add user input message to history

Sma1lboy · Oct 26, 2024 · da1463b · da1463b
1 parent a6fec39
commit da1463b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/backend/src/chat/chat.resolver.ts b/backend/src/chat/chat.resolver.ts
@@ -3,7 +3,7 @@ import { ChatCompletionChunk } from './chat.model';
 import { ChatProxyService, ChatService } from './chat.service';
 import { UserService } from 'src/user/user.service';
 import { Chat } from './chat.model';
-import { Message } from 'src/chat/message.model';
+import { Message, Role } from 'src/chat/message.model';
 import {
   NewChatInput,
   UpateChatTitleInput,
@@ -27,13 +27,15 @@ export class ChatResolver {
   })
   async *chatStream(@Args('input') input: ChatInput) {
     const iterator = this.chatProxyService.streamChat(input.message);
+    this.chatService.saveMessage(input.chatId, null, input.message, Role.User);
     try {
       for await (const chunk of iterator) {
         if (chunk) {
           await this.chatService.saveMessage(
             input.chatId,
             chunk.id,
             chunk.choices[0].delta.content,
+            Role.Model,
           );
           yield chunk;
         }
@@ -58,6 +60,7 @@ export class ChatResolver {
     return this.chatService.getMessageById(messageId);
   }
 
+  @UseGuards(ChatGuard)
   @Query(() => [Message])
   async getChatHistory(@Args('chatId') chatId: string): Promise<Message[]> {
     return this.chatService.getChatHistory(chatId);
@@ -70,7 +73,7 @@ export class ChatResolver {
   }
 
   // @Query(() => [Message])
-  // getModelTags(@Args('chatId') chatId: string): Message[] {
+  // getAvailableModelTags(@Args('chatId') chatId: string): Message[] {
   //   return this.chatService.getChatHistory(chatId);
   // }
 

diff --git a/backend/src/chat/chat.service.ts b/backend/src/chat/chat.service.ts
@@ -225,18 +225,19 @@ export class ChatService {
 
   async saveMessage(
     chatId: string,
-    chunkId: string,
+    chunkId: string | null,
     messageContent: string,
+    role: Role,
   ): Promise<Message> {
     // Find the chat instance
     const chat = await this.chatRepository.findOne({ where: { id: chatId } });
     if (!chat) throw new Error('Chat not found');
 
     // Create a new message associated with the chat
     const message = this.messageRepository.create({
-      id: chunkId,
+      id: chunkId || null,
       content: messageContent,
-      role: Role.Model,
+      role: role,
       chat,
       createdAt: new Date(),
     });