fix: update GraphQL mutation input type and clean up comments and add…

backend/src/chat/chat.controller.ts

@@ -22,13 +22,6 @@ export class ChatController {
     @GetAuthToken() userId: string,
   ) {
     try {
-      // Save user's message first
-      await this.chatService.saveMessage(
-        chatDto.chatId,
-        chatDto.message,
-        MessageRole.User,
-      );
-
       if (chatDto.stream) {
         // Streaming response
         res.setHeader('Content-Type', 'text/event-stream');
@@ -39,6 +32,7 @@ export class ChatController {
           chatId: chatDto.chatId,
           message: chatDto.message,
           model: chatDto.model,
+          role: MessageRole.User,
         });
 
         let fullResponse = '';
@@ -51,13 +45,6 @@ export class ChatController {
           }
         }
 
-        // Save the complete message
-        await this.chatService.saveMessage(
-          chatDto.chatId,
-          fullResponse,
-          MessageRole.Assistant,
-        );
-
         res.write('data: [DONE]\n\n');
         res.end();
       } else {
@@ -66,15 +53,8 @@ export class ChatController {
           chatId: chatDto.chatId,
           message: chatDto.message,
           model: chatDto.model,
+          role: MessageRole.User,
         });
-
-        // Save the complete message
-        await this.chatService.saveMessage(
-          chatDto.chatId,
-          response,
-          MessageRole.Assistant,
-        );
-
         res.json({ content: response });
       }
     } catch (error) {

backend/src/chat/chat.model.ts

@@ -75,11 +75,11 @@ class ChatCompletionDelta {
 
 @ObjectType('ChatCompletionChoiceType')
 class ChatCompletionChoice {
-  @Field()
-  index: number;
+  @Field({ nullable: true })
+  index: number | null;
 
-  @Field(() => ChatCompletionDelta)
-  delta: ChatCompletionDelta;
+  @Field(() => ChatCompletionDelta, { nullable: true })
+  delta: ChatCompletionDelta | null;
 
   @Field({ nullable: true })
   finishReason: string | null;
@@ -90,14 +90,14 @@ export class ChatCompletionChunk {
   @Field()
   id: string;
 
-  @Field()
-  object: string;
+  @Field({ nullable: true })
+  object: string | null;
 
-  @Field()
-  created: number;
+  @Field({ nullable: true })
+  created: number | null;
 
-  @Field()
-  model: string;
+  @Field({ nullable: true })
+  model: string | null;
 
   @Field({ nullable: true })
   systemFingerprint: string | null;

backend/src/chat/chat.resolver.ts

@@ -1,8 +1,8 @@
 import { Resolver, Subscription, Args, Query, Mutation } from '@nestjs/graphql';
-import { Chat, ChatCompletionChunk } from './chat.model';
+import { Chat, ChatCompletionChunk, StreamStatus } from './chat.model';
 import { ChatProxyService, ChatService } from './chat.service';
 import { UserService } from 'src/user/user.service';
-import { Message, MessageRole } from './message.model';
+import { Message } from './message.model';
 import {
   ChatInput,
   NewChatInput,
@@ -12,6 +12,7 @@ import { GetUserIdFromToken } from 'src/decorator/get-auth-token.decorator';
 import { Inject, Logger } from '@nestjs/common';
 import { JWTAuth } from 'src/decorator/jwt-auth.decorator';
 import { PubSubEngine } from 'graphql-subscriptions';
+import { Project } from 'src/project/project.model';
 @Resolver('Chat')
 export class ChatResolver {
   private readonly logger = new Logger('ChatResolver');
@@ -31,45 +32,65 @@ export class ChatResolver {
     resolve: (payload) => payload.chatStream,
   })
   async chatStream(@Args('input') input: ChatInput) {
-    return this.pubSub.asyncIterator(`chat_stream_${input.chatId}`);
+    const asyncIterator = this.pubSub.asyncIterator(
+      `chat_stream_${input.chatId}`,
+    );
+    return asyncIterator;
   }
-
   @Mutation(() => Boolean)
   @JWTAuth()
-  async triggerChatStream(@Args('input') input: ChatInput): Promise<boolean> {
+  async saveMessage(@Args('input') input: ChatInput): Promise<boolean> {
     try {
       await this.chatService.saveMessage(
         input.chatId,
         input.message,
-        MessageRole.User,
+        input.role,
       );
-
+      return true;
+    } catch (error) {
+      this.logger.error('Error in saveMessage:', error);
+      throw error;
+    }
+  }
+  @Mutation(() => Boolean)
+  @JWTAuth()
+  async triggerChatStream(@Args('input') input: ChatInput): Promise<boolean> {
+    try {
       const iterator = this.chatProxyService.streamChat(input);
       let accumulatedContent = '';
 
-      for await (const chunk of iterator) {
-        if (chunk) {
-          const enhancedChunk = {
-            ...chunk,
-            chatId: input.chatId,
-          };
+      try {
+        for await (const chunk of iterator) {
+          console.log('received chunk:', chunk);
+          if (chunk) {
+            const enhancedChunk = {
+              ...chunk,
+              chatId: input.chatId,
+            };
+
+            await this.pubSub.publish(`chat_stream_${input.chatId}`, {
+              chatStream: enhancedChunk,
+            });
+
+            if (chunk.choices?.[0]?.delta?.content) {
+              accumulatedContent += chunk.choices[0].delta.content;
+            }
+          }
+        }
+      } finally {
+        const finalChunk = await iterator.return();
+        console.log('finalChunk:', finalChunk);
 
+        if (finalChunk.value?.status === StreamStatus.DONE) {
           await this.pubSub.publish(`chat_stream_${input.chatId}`, {
-            chatStream: enhancedChunk,
+            chatStream: {
+              ...finalChunk.value,
+              chatId: input.chatId,
+            },
           });
-
-          if (chunk.choices[0]?.delta?.content) {
-            accumulatedContent += chunk.choices[0].delta.content;
-          }
         }
       }
 
-      await this.chatService.saveMessage(
-        input.chatId,
-        accumulatedContent,
-        MessageRole.Assistant,
-      );
-
       return true;
     } catch (error) {
       this.logger.error('Error in triggerChatStream:', error);
@@ -108,6 +129,19 @@ export class ChatResolver {
     return this.chatService.getChatDetails(chatId);
   }
 
+  @JWTAuth()
+  @Query(() => Project, { nullable: true })
+  async getCurProject(@Args('chatId') chatId: string): Promise<Project> {
+    try {
+      const response = await this.chatService.getProjectByChatId(chatId);
+      this.logger.log('Loaded project:', response);
+      return response;
+    } catch (error) {
+      this.logger.error('Failed to fetch project:', error);
+      throw new Error('Failed to fetch project');
+    }
+  }
+
   @Mutation(() => Chat)
   @JWTAuth()
   async createChat(

backend/src/chat/chat.service.ts

@@ -11,6 +11,7 @@ import {
 } from 'src/chat/dto/chat.input';
 import { CustomAsyncIterableIterator } from 'src/common/model-provider/types';
 import { OpenAIModelProvider } from 'src/common/model-provider/openai-model-provider';
+import { Project } from 'src/project/project.model';
 
 @Injectable()
 export class ChatProxyService {
@@ -98,6 +99,15 @@ export class ChatService {
     return chat;
   }
 
+  async getProjectByChatId(chatId: string): Promise<Project> {
+    const chat = await this.chatRepository.findOne({
+      where: { id: chatId, isDeleted: false },
+      relations: ['project'],
+    });
+
+    return chat ? chat.project : null;
+  }
+
   async createChat(userId: string, newChatInput: NewChatInput): Promise<Chat> {
     const user = await this.userRepository.findOne({ where: { id: userId } });
     if (!user) {

backend/src/chat/dto/chat.input.ts

@@ -1,5 +1,6 @@
 // DTOs for Project APIs
 import { InputType, Field } from '@nestjs/graphql';
+import { MessageRole } from '../message.model';
 
 @InputType()
 export class NewChatInput {
@@ -26,4 +27,6 @@ export class ChatInput {
 
   @Field()
   model: string;
+  @Field()
+  role: MessageRole;
 }

backend/src/common/model-provider/openai-model-provider.ts

@@ -118,7 +118,7 @@ export class OpenAIModelProvider implements IModelProvider {
     let streamIterator: AsyncIterator<OpenAIChatCompletionChunk> | null = null;
     const modelName = model || input.model;
     const queue = this.getQueueForModel(modelName);
-
+    let oldStreamValue: OpenAIChatCompletionChunk | null = null;
     const createStream = async () => {
       if (!stream) {
         const result = await queue.add(async () => {
@@ -145,6 +145,9 @@ export class OpenAIModelProvider implements IModelProvider {
           const currentIterator = await createStream();
           const chunk = await currentIterator.next();
           const chunkValue = chunk.value as OpenAIChatCompletionChunk;
+          console.log('isDone:', chunk.done);
+          console.log('chunk:', chunk);
+          if (!chunk.done) oldStreamValue = chunkValue;
           return {
             done: chunk.done,
             value: {
@@ -159,9 +162,23 @@ export class OpenAIModelProvider implements IModelProvider {
         }
       },
       async return() {
+        console.log(stream);
+        console.log(streamIterator);
+        console.log('return() called');
         stream = null;
         streamIterator = null;
-        return { done: true, value: undefined };
+        return {
+          done: true,
+          value: {
+            ...oldStreamValue,
+            status: StreamStatus.DONE,
+            choices: [
+              {
+                finishReason: 'stop',
+              },
+            ],
+          },
+        };
       },
       async throw(error) {
         stream = null;

frontend/src/api/ChatStreamAPI.ts

@@ -0,0 +1,71 @@
+import { ChatInputType } from '@/graphql/type';
+
+export const startChatStream = async (
+  input: ChatInputType,
+  token: string,
+  stream: boolean = false // Default to non-streaming for better performance
+): Promise<string> => {
+  if (!token) {
+    throw new Error('Not authenticated');
+  }
+  const { chatId, message, model } = input;
+  const response = await fetch('/api/chat', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${token}`,
+    },
+    body: JSON.stringify({
+      chatId,
+      message,
+      model,
+      stream,
+    }),
+  });
+
+  if (!response.ok) {
+    throw new Error(
+      `Network response was not ok: ${response.status} ${response.statusText}`
+    );
+  }
+  // TODO: Handle streaming responses properly
+  // if (stream) {
+  //   // For streaming responses, aggregate the streamed content
+  //   let fullContent = '';
+  //   const reader = response.body?.getReader();
+  //   if (!reader) {
+  //     throw new Error('No reader available');
+  //   }
+
+  //   while (true) {
+  //     const { done, value } = await reader.read();
+  //     if (done) break;
+
+  //     const text = new TextDecoder().decode(value);
+  //     const lines = text.split('\n\n');
+
+  //     for (const line of lines) {
+  //       if (line.startsWith('data: ')) {
+  //         const data = line.slice(5);
+  //         if (data === '[DONE]') break;
+  //         try {
+  //           const { content } = JSON.parse(data);
+  //           if (content) {
+  //             fullContent += content;
+  //           }
+  //         } catch (e) {
+  //           console.error('Error parsing SSE data:', e);
+  //         }
+  //       }
+  //     }
+  //   }
+  //   return fullContent;
+  // } else {
+  //   // For non-streaming responses, return the content directly
+  //   const data = await response.json();
+  //   return data.content;
+  // }
+
+  const data = await response.json();
+  return data.content;
+};