feat: langfuse tracing with user id / session id / chat engine tag (#149

) * feat: tracing with user id and tag * add sessionId
pingcap · Jun 3, 2024 · b7ef630 · b7ef630
1 parent 3fcb6c2
commit b7ef630
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 10 deletions.
diff --git a/src/app/api/v1/chats/route.ts b/src/app/api/v1/chats/route.ts
@@ -41,7 +41,7 @@ export const POST = defineHandler({
     messages,
   } = body;
 
-  const [engine, engineOptions] = await getChatEngineConfig(body.engine);
+  const [engineId, engine, engineOptions] = await getChatEngineConfig(body.engine);
 
   // TODO: need refactor, it is too complex now
   // For chat page, create a chat and return the session ID (url_key) first.
@@ -59,6 +59,7 @@ export const POST = defineHandler({
 
     return await createChat({
       engine,
+      engine_id: engineId,
       engine_options: JSON.stringify(engineOptions),
       created_at: new Date(),
       created_by: userId,

diff --git a/src/app/api/v1/indexes/[name]/retrieve/route.ts b/src/app/api/v1/indexes/[name]/retrieve/route.ts
@@ -26,7 +26,7 @@ export const POST = defineHandler({
     notFound();
   }
 
-  const [engine, engineOptions] = await getChatEngineConfig(body.engine);
+  const [engineId, engine, engineOptions] = await getChatEngineConfig(body.engine);
   const {
     llm: llmConfig = {
       provider: LLMProvider.OPENAI,

diff --git a/src/core/db/schema.d.ts b/src/core/db/schema.d.ts
@@ -16,6 +16,11 @@ export type JsonPrimitive = boolean | number | string | null;
 
 export type JsonValue = JsonArray | JsonObject | JsonPrimitive;
 
+export interface App {
+  app_id: string | null;
+  name: string | null;
+}
+
 export interface AppAccessToken {
   app_id: string;
   token: string;
@@ -33,6 +38,7 @@ export interface Chat {
   deleted_at: Date | null;
   deleted_by: string | null;
   engine: string;
+  engine_id: number | null;
   engine_options: Json;
   id: Generated<number>;
   title: string;
@@ -206,6 +212,7 @@ export interface Status {
 }
 
 export interface DB {
+  app: App;
   app_access_token: AppAccessToken;
   authentication_provider: AuthenticationProvider;
   chat: Chat;

diff --git a/src/core/repositories/chat_engine.ts b/src/core/repositories/chat_engine.ts
@@ -24,6 +24,16 @@ export async function getChatEngine (id: number) {
     .executeTakeFirst();
 }
 
+export async function getChatEngineNameByID (id?: number | null) {
+  if (!id) return undefined;
+  const res = await getDb()
+    .selectFrom('chat_engine')
+    .select(['name'])
+    .where('id', '=', id)
+    .executeTakeFirst();
+  return res?.name;
+}
+
 export async function getDefaultChatEngine () {
   return await getDb()
     .selectFrom('chat_engine')
@@ -33,16 +43,16 @@ export async function getDefaultChatEngine () {
     .executeTakeFirstOrThrow();
 }
 
-export async function getChatEngineConfig (engineConfigId?: number): Promise<[string, ChatEngineOptions]> {
+export async function getChatEngineConfig (engineConfigId?: number): Promise<[number, string, ChatEngineOptions]> {
   if (engineConfigId) {
-    const chatEngine = await getChatEngine(engineConfigId);
-    if (!chatEngine) {
+    const config = await getChatEngine(engineConfigId);
+    if (!config) {
       throw CHAT_ENGINE_NOT_FOUND_ERROR.format(engineConfigId);
     }
-    return [chatEngine.engine, chatEngine.engine_options];
+    return [config.id, config.engine, config.engine_options];
   } else {
     const config = await getDefaultChatEngine();
-    return [config.engine, config.engine_options];
+    return [config.id, config.engine, config.engine_options];
   }
 }
 

diff --git a/src/core/services/llamaindex/chating.ts b/src/core/services/llamaindex/chating.ts
@@ -1,6 +1,6 @@
 import {getDb} from '@/core/db';
 import {type Chat, listChatMessages, updateChatMessage} from '@/core/repositories/chat';
-import {ChatEngineRequiredOptions} from '@/core/repositories/chat_engine';
+import {ChatEngineRequiredOptions, getChatEngineNameByID} from '@/core/repositories/chat_engine';
 import {getDocumentsBySourceUris} from "@/core/repositories/document";
 import {GraphRetrieverSearchOptions} from "@/core/schema/chat_engines/condense_question";
 import {AppChatService, type ChatOptions, type ChatStreamEvent} from '@/core/services/chating';
@@ -13,15 +13,14 @@ import {
   Entity,
   KnowledgeGraphClient,
   Relationship,
-  SearchOptions,
   SearchResult
 } from "@/lib/knowledge-graph/client";
 import {uuidToBin} from '@/lib/kysely';
 import {buildEmbedding} from '@/lib/llamaindex/builders/embedding';
 import {buildLLM} from "@/lib/llamaindex/builders/llm";
 import {buildReranker} from "@/lib/llamaindex/builders/reranker";
 import {LLMConfig, LLMProvider} from "@/lib/llamaindex/config/llm";
-import { type RerankerConfig, RerankerProvider } from '@/lib/llamaindex/config/reranker';
+import { type RerankerConfig } from '@/lib/llamaindex/config/reranker';
 import {ManagedAsyncIterable} from '@/lib/ManagedAsyncIterable';
 import {LangfuseTraceClient} from "langfuse";
 import {Liquid} from 'liquidjs';
@@ -106,13 +105,20 @@ export class LlamaindexChatService extends AppChatService {
       prompts
     } = engineOptions;
 
+    const chatEngineName = await getChatEngineNameByID(chat.engine_id);
+
     // Init tracing.
     const trace = this.langfuse?.trace({
       name: 'chatting',
       input: {
         history: options.history,
         userInput: options.userInput
       },
+      userId: chat.created_by,
+      sessionId: chat.url_key,
+      tags: [
+        `chat_engine:${chatEngineName}`
+      ],
       metadata: {
         chat_id: chat.id,
         chat_slug: chat.url_key,