feat: support synopsis entities (#154)

* refactor: return trace id instead trace object * remove unused imports * fix: record chat url instead of chat url key * feat: support synopsis entities * add default site url
pingcap · Jun 5, 2024 · e46f3ca · e46f3ca
1 parent bd2fd8b
commit e46f3ca
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 14 deletions.
diff --git a/src/app/api/v1/chat_engines/[id]/test/route.ts b/src/app/api/v1/chat_engines/[id]/test/route.ts
@@ -62,8 +62,10 @@ export const POST = defineHandler({
 
     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, {
+    if (chatResult.traceId) {
+      await item.link(langfuse.trace({
+        id: chatResult.traceId,
+      }), run.name, {
         description: run.description,
         metadata: run.metadata
       });

diff --git a/src/app/api/v1/chats/route.ts b/src/app/api/v1/chats/route.ts
@@ -1,7 +1,6 @@
 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';
@@ -113,8 +112,7 @@ export const POST = defineHandler({
   if (body.stream) {
     return chatResult.toResponse();
   } else {
-    const { trace, ...result} = chatResult as unknown as ChatNonStreamingResult;
-    return result;
+    return chatResult;
   }
 });
 

diff --git a/src/core/services/chating.ts b/src/core/services/chating.ts
@@ -17,7 +17,7 @@ export type ChatOptions = {
 export type ChatStreamEvent = {
   status: AppChatStreamState;
   statusMessage: string;
-  trace?: LangfuseTraceClient;
+  traceId?: string;
   traceURL?: string;
   sources: AppChatStreamSource[];
   content: string;
@@ -26,7 +26,7 @@ export type ChatStreamEvent = {
 }
 
 export interface ChatNonStreamingResult {
-  trace?: LangfuseTraceClient;
+  traceId?: string;
   traceURL: string;
   content: string;
   sources: AppChatStreamSource[];
@@ -82,8 +82,8 @@ export abstract class AppChatService extends AppIndexBaseService {
           if (chunk.traceURL && chunk.traceURL.length > 0) {
             chatResult.traceURL = chunk.traceURL;
           }
-          if (chunk.trace) {
-            chatResult.trace = chunk.trace;
+          if (chunk.traceId && chunk.traceId.length > 0) {
+            chatResult.traceId = chunk.traceId;
           }
         }
         chatResult.state = AppChatStreamState.FINISHED;

diff --git a/src/core/services/llamaindex/chating.ts b/src/core/services/llamaindex/chating.ts
@@ -121,7 +121,7 @@ export class LlamaindexChatService extends AppChatService {
       ],
       metadata: {
         chat_id: chat.id,
-        chat_slug: chat.url_key,
+        chat_url: `${process.env.SITE_URL || 'https://tidb.ai'}/c/${chat.url_key}`,
         chat_engine_type: chat.engine,
         chat_engine_options: engineOptions,
       },
@@ -130,7 +130,7 @@ export class LlamaindexChatService extends AppChatService {
     yield {
       status: AppChatStreamState.CREATING,
       sources: [],
-      trace: trace,
+      traceId: trace?.traceId,
       traceURL: trace?.getTraceUrl(),
       statusMessage: '',
       retrieveId: undefined,
@@ -200,8 +200,15 @@ export class LlamaindexChatService extends AppChatService {
         );
 
         // Flatten relationships and entities.
+
+        // Notice: Synopsis entities do not link to any relationship, so we need to append them to the entity list.
+        const synopsisEntities = result.entities.filter(entity => entity.entity_type === 'synopsis');
+        const rerankedEntities = result.document_relationships.map(dr => dr.entities).flat().filter(deduplicateItems('id'));
+        result.entities = [
+          ...synopsisEntities,
+          ...rerankedEntities
+        ];
         result.relationships = result.document_relationships.map(dr => dr.relationships).flat().filter(deduplicateItems('id'));
-        result.entities = result.document_relationships.map(dr => dr.entities).flat().filter(deduplicateItems('id'));
       }
 
       kgContext = result;
@@ -401,7 +408,7 @@ export class LlamaindexChatService extends AppChatService {
     const duration = DateTime.now().diff(start, 'milliseconds').milliseconds;
 
     kgSearchSpan?.end({
-      output: searchResult,
+      output: JSON.parse(JSON.stringify(searchResult)),
     });
     console.log(`[KG-Retrieving] Finish knowledge graph searching, take ${duration} ms.`);
     // Fixme: DO NOT MUTATE THE LANGFUSE OUTPUT. Langfuse always delays it's push operation, mutations before pushing will affect langfuse tracking.

diff --git a/src/core/services/llamaindex/retrieving.ts b/src/core/services/llamaindex/retrieving.ts
@@ -276,7 +276,8 @@ export class LlamaindexRetrieverWrapper implements BaseRetriever {
           metadata: {
             //// MARK: we don't need the metadata from extractors, they are for embedding.
             // ...chunk.metadata,
-            sourceUri: chunk.document_uri || url,
+            source_title: chunk.document_name,
+            source_uri: chunk.document_uri || url,
             ...restMetadata,
           },
           relationships: Object.fromEntries(Object.entries(chunk.relationships).map(([k, v]) => {

diff --git a/src/lib/knowledge-graph/client.ts b/src/lib/knowledge-graph/client.ts
@@ -12,6 +12,7 @@ export interface Entity {
   name: string;
   description: string;
   meta: Record<string, any> | null;
+  entity_type: 'original' | 'synopsis'
 }
 
 export interface Relationship {