Merge pull request #82 from import-ai/fix/message_status

Update message status when error

Author: LucienShui

.github/workflows/ci.yaml

@@ -29,7 +29,7 @@ jobs:
       - name: Check if dependencies changed
         id: depcheck
         run: |
-          if git diff --name-only origin/main | grep -E 'package.json|pnpm-lock.yaml|pnpm-workspace.yaml'; then
+          if git diff --name-only origin/main | grep -E 'package.json|pnpm-lock.yaml|pnpm-workspace.yaml|runtime.Dockerfile'; then
             echo "changed=true" >> $GITHUB_OUTPUT
           else
             echo "changed=false" >> $GITHUB_OUTPUT

src/messages/messages.service.ts

@@ -9,6 +9,7 @@ import {
 import { CreateMessageDto } from 'src/messages/dto/create-message.dto';
 import { User } from 'src/user/user.entity';
 import { SearchService } from 'src/search/search.service';
+import { ChatDeltaResponse } from '../wizard/dto/chat-response.dto';
 
 @Injectable()
 export class MessagesService {
@@ -74,11 +75,9 @@ export class MessagesService {
     return delta ? (source || '') + delta : source;
   }
 
-  async updateOpenAIMessage(
-    id: string,
-    deltaMessage: Partial<OpenAIMessage>,
-    attrs?: Record<string, any>,
-  ) {
+  async updateDelta(id: string, delta: ChatDeltaResponse) {
+    const deltaMessage: Partial<OpenAIMessage> = delta.message;
+
     const message = await this.messageRepository.findOneOrFail({
       where: { id },
     });
@@ -100,9 +99,9 @@ export class MessagesService {
     }
     // <<< OpenAI Message
     message.status = MessageStatus.STREAMING;
-    if (attrs) {
+    if (delta.attrs) {
       message.attrs = message.attrs || {};
-      Object.assign(message.attrs, attrs);
+      Object.assign(message.attrs, delta.attrs);
     }
     return await this.messageRepository.save(message);
   }

src/wizard/stream.service.ts

@@ -109,10 +109,9 @@ export class StreamService {
         if (!context.messageId) {
           throw new Error('Message ID is not set in context');
         }
-        const message: Message = await this.messagesService.updateOpenAIMessage(
+        const message: Message = await this.messagesService.updateDelta(
           context.messageId,
-          chunk.message,
-          chunk.attrs,
+          chunk,
         );
 
         context.message = message.message;
@@ -133,6 +132,18 @@ export class StreamService {
       } else if (chunk.response_type === 'done') {
         // Do nothing, this is the end of the stream
       } else if (chunk.response_type === 'error') {
+        if (context.messageId) {
+          await this.messagesService.update(
+            context.messageId,
+            namespaceId,
+            conversationId,
+            {
+              status: MessageStatus.FAILED,
+            },
+            true,
+          );
+        }
+
         const err = new Error(chunk.message || 'Unknown error');
         err.name = 'AgentError';
         throw err;
@@ -145,17 +156,6 @@ export class StreamService {
     };
   }
 
-  chatStream(body: Record<string, any>): Observable<MessageEvent> {
-    return new Observable<MessageEvent>((subscriber) => {
-      this.stream('/api/v1/wizard/stream', body, (data) => {
-        subscriber.next({ data });
-        return Promise.resolve();
-      })
-        .then(() => subscriber.complete())
-        .catch((err) => subscriber.error(err));
-    });
-  }
-
   findOneOrFail(messages: Message[], messageId: string): Message {
     const message = messages.find((m) => m.id === messageId);
     if (!message) {

src/wizard/wizard.controller.ts

@@ -16,12 +16,6 @@ export class WizardController {
     return await this.wizardService.collect(req.user, data);
   }
 
-  @Post('chat/stream')
-  @Sse()
-  chat(@Body() body: Record<string, any>) {
-    return this.wizardService.streamService.chatStream(body);
-  }
-
   @Post('ask')
   @Sse()
   async ask(@Req() req, @Body() body: AgentRequestDto) {