fix: save history messages when create non-session chat (#164)

src/app/api/v1/chats/route.ts

@@ -1,14 +1,17 @@
-import {type Chat, createChat, getChatByUrlKey, listChats} from '@/core/repositories/chat';
+import {type Chat, createChat, createChatMessages, getChatByUrlKey, listChats} from '@/core/repositories/chat';
 import {getChatEngineByIdOrName} from '@/core/repositories/chat_engine';
 import {getIndexByNameOrThrow} from '@/core/repositories/index_';
 import {LlamaindexChatService} from '@/core/services/llamaindex/chating';
 import {toPageRequest} from '@/lib/database';
-import {CHAT_CAN_NOT_ASSIGN_SESSION_ID_ERROR} from '@/lib/errors';
+import {
+  CHAT_CAN_NOT_ASSIGN_SESSION_ID_ERROR,
+  CHAT_FAILED_TO_CREATE_ERROR,
+  CHAT_NOT_FOUND_ERROR
+} from '@/lib/errors';
 import {defineHandler} from '@/lib/next/handler';
 import {baseRegistry} from '@/rag-spec/base';
 import {getFlow} from '@/rag-spec/createFlow';
 import {Langfuse} from "langfuse";
-import {notFound} from 'next/navigation';
 import {NextResponse} from 'next/server';
 import {z} from 'zod';
 
@@ -69,24 +72,39 @@ export const POST = defineHandler({
 
   const lastUserMessage = messages.findLast(m => m.role === 'user')?.content ?? '';
 
-  // For Ask Widget.
-  let chat: Chat | undefined;
+  // For Ask Widget / API.
+  let chat: Chat;
   let sessionId = body.sessionId;
   if (!sessionId) {
-    chat = await createChat({
+    chat = (await createChat({
       engine: engine.engine,
       engine_id: engine.id,
       engine_name: engine.name,
       engine_options: JSON.stringify(engine.engine_options),
       created_at: new Date(),
       created_by: userId,
       title: limitTitleLength(body.name ?? lastUserMessage ?? DEFAULT_CHAT_TITLE),
-    });
+    }))!;
+    if (!chat) {
+      throw CHAT_FAILED_TO_CREATE_ERROR;
+    }
     sessionId = chat.url_key;
+    const previousMessages = messages.length > 1 ? messages.slice(0, -1) : [];
+    if (previousMessages.length > 0) {
+      await createChatMessages(previousMessages.map((m, idx) => ({
+        chat_id: chat.id,
+        content: m.content,
+        role: m.role,
+        ordinal: idx,
+        status: 'SUCCEED',
+        options: '{}',
+        created_at: new Date(),
+      })));
+    }
   } else {
-    chat = await getChatByUrlKey(sessionId);
+    chat = (await getChatByUrlKey(sessionId))!;
     if (!chat) {
-      notFound();
+      throw CHAT_NOT_FOUND_ERROR.format(sessionId);
     }
   }
 

src/core/repositories/chat.ts

@@ -93,6 +93,13 @@ export async function createChatMessage (create: CreateChatMessage) {
   return (await getChatMessage(Number(insertId)))!;
 }
 
+export async function createChatMessages (create: CreateChatMessage[]) {
+  return await getDb()
+    .insertInto('chat_message')
+    .values(create)
+    .executeTakeFirstOrThrow();
+}
+
 export async function updateChatMessage (id: number, update: UpdateChatMessage) {
   await getDb()
     .updateTable('chat_message')

src/core/services/chating.ts

@@ -9,6 +9,7 @@ import { notFound } from 'next/navigation';
 export type ChatOptions = {
   userInput: string;
   userId: string;
+  requestMessage: ChatMessage;
   respondMessage: ChatMessage;
   history: ChatMessage[];
 }
@@ -39,14 +40,20 @@ export abstract class AppChatService extends AppIndexBaseService {
   chat(sessionId: string, userId: string, userInput: string, regenerating: boolean, stream: false): Promise<ChatNonStreamingResult>
   async chat(sessionId: string, userId: string, userInput: string, regenerating: boolean, stream: true | false): Promise<AppChatStream | ChatNonStreamingResult> {
     const { chat, history } = await this.getSessionInfo(sessionId, userId);
-    const respondMessage = await this.startChat(chat, history, userInput, regenerating);
+    const { requestMessage, respondMessage } = await this.startChat(chat, history, userInput, regenerating);
 
     if (stream) {
       return new AppChatStream(sessionId, respondMessage.id, async controller => {
         try {
           let content = '';
           let retrieveIds = new Set<number>();
-          for await (const chunk of this.run(chat, { userInput, history, userId, respondMessage })) {
+          for await (const chunk of this.run(chat, {
+            userInput,
+            history,
+            userId,
+            requestMessage,
+            respondMessage
+          })) {
             controller.appendText(chunk.content, chunk.status === AppChatStreamState.CREATING /* force sends an empty text chunk first, to avoid a dependency BUG */);
             controller.setChatState(chunk.status, chunk.statusMessage);
             controller.setTraceURL(chunk.traceURL);
@@ -73,7 +80,13 @@ export abstract class AppChatService extends AppIndexBaseService {
       };
       try {
         let retrieveIds = new Set<number>();
-        for await (const chunk of this.run(chat, { userInput, history, userId, respondMessage })) {
+        for await (const chunk of this.run(chat, {
+          userInput,
+          history,
+          userId,
+          requestMessage,
+          respondMessage
+        })) {
           chatResult.content += chunk.content;
           chatResult.sources = chunk.sources;
           if (chunk.retrieveId) {
@@ -131,11 +144,15 @@ export abstract class AppChatService extends AppIndexBaseService {
     });
   }
 
-  private async startChat (chat: Chat, history: ChatMessage[], userInput: string, regenerating: boolean) {
+  private async startChat (chat: Chat, history: ChatMessage[], userInput: string, regenerating: boolean): Promise<{
+    requestMessage: ChatMessage,
+    respondMessage: ChatMessage,
+  }> {
     return await tx(async () => {
       let ordinal = history.length;
+      let requestMessage = history[ordinal - 1];
       if (!regenerating) {
-        await createChatMessage({
+        requestMessage = await createChatMessage({
           role: 'user',
           chat_id: chat.id,
           content: userInput,
@@ -145,7 +162,7 @@ export abstract class AppChatService extends AppIndexBaseService {
           options: JSON.stringify({}),
         });
       }
-      return await createChatMessage({
+      let respondMessage = await createChatMessage({
         role: 'assistant',
         chat_id: chat.id,
         content: '',
@@ -154,6 +171,7 @@ export abstract class AppChatService extends AppIndexBaseService {
         ordinal: ordinal,
         options: JSON.stringify({}),
       });
+      return { requestMessage, respondMessage };
     });
   }
 

src/core/services/llamaindex/chating.ts

@@ -124,6 +124,8 @@ export class LlamaindexChatService extends AppChatService {
         chat_url: `${process.env.SITE_URL || 'https://tidb.ai'}/c/${chat.url_key}`,
         chat_engine_type: chat.engine,
         chat_engine_options: engineOptions,
+        request_message_id: options.requestMessage.id,
+        respond_message_id: options.respondMessage.id,
       },
     });
 

src/lib/errors/api_errors.ts

@@ -102,3 +102,7 @@ export const CHAT_ENGINE_NOT_FOUND_ERROR = APIError.new(`Specified chat engine %
  * Chat Related
  */
 export const CHAT_CAN_NOT_ASSIGN_SESSION_ID_ERROR = APIError.new( 'Cannot assign sessionId when creating chats.', 400);
+
+export const CHAT_NOT_FOUND_ERROR = APIError.new(`Chat <%s> not found`, 404);
+
+export const CHAT_FAILED_TO_CREATE_ERROR = APIError.new(`Failed to create chat`, 500);