diff --git a/src/app/api/v1/chats/route.ts b/src/app/api/v1/chats/route.ts index 33cb61f9..bdc67c5b 100644 --- a/src/app/api/v1/chats/route.ts +++ b/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,11 +72,11 @@ 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, @@ -81,12 +84,27 @@ export const POST = defineHandler({ 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); } } diff --git a/src/core/repositories/chat.ts b/src/core/repositories/chat.ts index ed57f035..434b4229 100644 --- a/src/core/repositories/chat.ts +++ b/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') diff --git a/src/core/services/chating.ts b/src/core/services/chating.ts index 649a4dd9..9c233792 100644 --- a/src/core/services/chating.ts +++ b/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 async chat(sessionId: string, userId: string, userInput: string, regenerating: boolean, stream: true | false): Promise { 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(); - 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(); - 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 }; }); } diff --git a/src/core/services/llamaindex/chating.ts b/src/core/services/llamaindex/chating.ts index 530cbec7..01bfdd01 100644 --- a/src/core/services/llamaindex/chating.ts +++ b/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, }, }); diff --git a/src/lib/errors/api_errors.ts b/src/lib/errors/api_errors.ts index a8ab42c2..d498456c 100644 --- a/src/lib/errors/api_errors.ts +++ b/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);