feat(messageHandler): fix duplicate messageIDs

This commit is contained in:
ItzCrazyKns 2024-10-17 18:50:43 +05:30
parent 24e1919c5e
commit 5aad2febda
1 changed files with 30 additions and 17 deletions

View File

@ -10,8 +10,8 @@ import type { BaseChatModel } from '@langchain/core/language_models/chat_models'
import type { Embeddings } from '@langchain/core/embeddings'; import type { Embeddings } from '@langchain/core/embeddings';
import logger from '../utils/logger'; import logger from '../utils/logger';
import db from '../db'; import db from '../db';
import { chats, messages } from '../db/schema'; import { chats, messages as messagesSchema } from '../db/schema';
import { eq } from 'drizzle-orm'; import { eq, asc, gt } from 'drizzle-orm';
import crypto from 'crypto'; import crypto from 'crypto';
type Message = { type Message = {
@ -71,7 +71,7 @@ const handleEmitterEvents = (
emitter.on('end', () => { emitter.on('end', () => {
ws.send(JSON.stringify({ type: 'messageEnd', messageId: messageId })); ws.send(JSON.stringify({ type: 'messageEnd', messageId: messageId }));
db.insert(messages) db.insert(messagesSchema)
.values({ .values({
content: recievedMessage, content: recievedMessage,
chatId: chatId, chatId: chatId,
@ -106,7 +106,9 @@ export const handleMessage = async (
const parsedWSMessage = JSON.parse(message) as WSMessage; const parsedWSMessage = JSON.parse(message) as WSMessage;
const parsedMessage = parsedWSMessage.message; const parsedMessage = parsedWSMessage.message;
const id = crypto.randomBytes(7).toString('hex'); const humanMessageId =
parsedMessage.messageId ?? crypto.randomBytes(7).toString('hex');
const aiMessageId = crypto.randomBytes(7).toString('hex');
if (!parsedMessage.content) if (!parsedMessage.content)
return ws.send( return ws.send(
@ -141,7 +143,7 @@ export const handleMessage = async (
parsedWSMessage.optimizationMode, parsedWSMessage.optimizationMode,
); );
handleEmitterEvents(emitter, ws, id, parsedMessage.chatId); handleEmitterEvents(emitter, ws, aiMessageId, parsedMessage.chatId);
const chat = await db.query.chats.findFirst({ const chat = await db.query.chats.findFirst({
where: eq(chats.id, parsedMessage.chatId), where: eq(chats.id, parsedMessage.chatId),
@ -159,18 +161,29 @@ export const handleMessage = async (
.execute(); .execute();
} }
await db const messageExists = await db.query.messages.findFirst({
.insert(messages) where: eq(messagesSchema.messageId, humanMessageId),
.values({ });
content: parsedMessage.content,
chatId: parsedMessage.chatId, if (!messageExists) {
messageId: id, await db
role: 'user', .insert(messagesSchema)
metadata: JSON.stringify({ .values({
createdAt: new Date(), content: parsedMessage.content,
}), chatId: parsedMessage.chatId,
}) messageId: humanMessageId,
.execute(); role: 'user',
metadata: JSON.stringify({
createdAt: new Date(),
}),
})
.execute();
} else {
await db
.delete(messagesSchema)
.where(gt(messagesSchema.id, messageExists.id))
.execute();
}
} else { } else {
ws.send( ws.send(
JSON.stringify({ JSON.stringify({