feat: support test chat engine with dataset (#152)

package.json

@@ -47,6 +47,7 @@
     "@types/luxon": "^3.4.2",
     "@types/mdx": "^2.0.11",
     "@vercel/blob": "^0.19.0",
+    "async": "^3.2.5",
     "change-case": "^5.4.2",
     "cheerio": "1.0.0-rc.11",
     "class-variance-authority": "^0.7.0",
@@ -112,6 +113,7 @@
     "@faker-js/faker": "^8.3.1",
     "@tailwindcss/typography": "^0.5.10",
     "@tanstack/table-core": "^8.11.7",
+    "@types/async": "^3.2.24",
     "@types/d3": "^7.4.3",
     "@types/hast": "^3.0.3",
     "@types/is-hotkey": "^0.1.10",

src/app/api/v1/chat_engines/[id]/test/route.ts

@@ -0,0 +1,75 @@
+import {createChat} 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 { defineHandler } from '@/lib/next/handler';
+import {baseRegistry} from "@/rag-spec/base";
+import {getFlow} from "@/rag-spec/createFlow";
+import {Langfuse} from "langfuse";
+import { z } from 'zod';
+import { eachOfLimit } from 'async';
+
+const SpecificChatEngineSchema = z.object({
+  id: z.string(),
+});
+
+const RunTestParamsSchema = z.object({
+  index_name: z.string().default('default'),
+  dataset_name: z.string(),
+  run: z.object({
+    name: z.string(),
+    description: z.string().optional(),
+    metadata: z.record(z.string()).optional(),
+  }),
+  concurrency: z.number().default(10)
+});
+
+export const POST = defineHandler({
+  auth: 'admin',
+  params: SpecificChatEngineSchema,
+  body: RunTestParamsSchema,
+}, async ({
+  params: { id },
+  body: {
+    dataset_name,
+    index_name,
+    concurrency,
+    run
+  },
+}) => {
+  const engine = await getChatEngineByIdOrName(id)
+
+  const langfuse = new Langfuse();
+  const dataset = await langfuse.getDataset(dataset_name);
+
+  const index = await getIndexByNameOrThrow(index_name);
+  const flow = await getFlow(baseRegistry);
+  const chatService = new LlamaindexChatService({ flow, index, langfuse });
+
+  await eachOfLimit(dataset.items, concurrency, async (item: any) => {
+    const userMessage = item.input;
+    const userId = 'tester';
+    const 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: userMessage,
+    });
+
+    console.log(`[Testing] Testing with question (chat engine: ${id}): ${userMessage}`)
+    const chatResult = await chatService.chat(chat.url_key, userId, userMessage, false, false);
+    if (chatResult.trace) {
+      await item.link(chatResult.trace, run.name, {
+        description: run.description,
+        metadata: run.metadata
+      });
+    }
+  });
+});
+
+export const dynamic = 'force-dynamic';
+
+export const maxDuration = 150;

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

@@ -1,6 +1,7 @@
 import {type Chat, createChat, getChatByUrlKey, listChats} from '@/core/repositories/chat';
 import {getChatEngineByIdOrName} from '@/core/repositories/chat_engine';
 import {getIndexByNameOrThrow} from '@/core/repositories/index_';
+import {ChatNonStreamingResult} from "@/core/services/chating";
 import {LlamaindexChatService} from '@/core/services/llamaindex/chating';
 import {toPageRequest} from '@/lib/database';
 import {CHAT_CAN_NOT_ASSIGN_SESSION_ID_ERROR} from '@/lib/errors';
@@ -107,12 +108,13 @@ export const POST = defineHandler({
   }
 
   const lastUserMessage = messages.findLast(m => m.role === 'user')?.content ?? '';
-  const chatResult = await chatService.chat(sessionId, userId, lastUserMessage, body.regenerate ?? false, body.stream);
+  const chatResult = await chatService.chat(sessionId, userId, lastUserMessage, body.regenerate ?? false, body.stream as any);
 
   if (body.stream) {
     return chatResult.toResponse();
   } else {
-    return chatResult;
+    const { trace, ...result} = chatResult as unknown as ChatNonStreamingResult;
+    return result;
   }
 });
 

src/core/services/chating.ts

@@ -17,6 +17,7 @@ export type ChatOptions = {
 export type ChatStreamEvent = {
   status: AppChatStreamState;
   statusMessage: string;
+  trace?: LangfuseTraceClient;
   traceURL?: string;
   sources: AppChatStreamSource[];
   content: string;
@@ -25,6 +26,7 @@ export type ChatStreamEvent = {
 }
 
 export interface ChatNonStreamingResult {
+  trace?: LangfuseTraceClient;
   traceURL: string;
   content: string;
   sources: AppChatStreamSource[];
@@ -33,7 +35,9 @@ export interface ChatNonStreamingResult {
 
 export abstract class AppChatService extends AppIndexBaseService {
 
-  async chat (sessionId: string, userId: string, userInput: string, regenerating: boolean, stream: boolean = true): Promise<any> {
+  chat(sessionId: string, userId: string, userInput: string, regenerating: boolean, stream: true): Promise<AppChatStream>;
+  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);
 
@@ -78,6 +82,9 @@ export abstract class AppChatService extends AppIndexBaseService {
           if (chunk.traceURL && chunk.traceURL.length > 0) {
             chatResult.traceURL = chunk.traceURL;
           }
+          if (chunk.trace) {
+            chatResult.trace = chunk.trace;
+          }
         }
         chatResult.state = AppChatStreamState.FINISHED;
         await this.finishChat(respondMessage, chatResult.content, retrieveIds);

src/core/services/llamaindex/chating.ts

@@ -130,6 +130,7 @@ export class LlamaindexChatService extends AppChatService {
     yield {
       status: AppChatStreamState.CREATING,
       sources: [],
+      trace: trace,
       traceURL: trace?.getTraceUrl(),
       statusMessage: '',
       retrieveId: undefined,
@@ -376,8 +377,6 @@ export class LlamaindexChatService extends AppChatService {
     trace?.update({
       output: finalResponse
     });
-
-    await this.langfuse?.flushAsync();
   }
 
   async searchKnowledgeGraph (