@lobehub/chat 1.79.4 → 1.79.6

package/src/server/modules/AgentRuntime/trace.ts

@@ -0,0 +1,107 @@
+import { after } from 'next/server';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import {
+  LOBE_CHAT_OBSERVATION_ID,
+  LOBE_CHAT_TRACE_ID,
+  TracePayload,
+  TraceTagMap,
+} from '@/const/trace';
+import { ChatStreamCallbacks, ChatStreamPayload } from '@/libs/agent-runtime';
+import { TraceClient } from '@/libs/traces';
+
+export interface AgentChatOptions {
+  enableTrace?: boolean;
+  provider: string;
+  trace?: TracePayload;
+}
+
+export const createTraceOptions = (
+  payload: ChatStreamPayload,
+  { trace: tracePayload, provider }: AgentChatOptions,
+) => {
+  const { messages, model, tools, ...parameters } = payload;
+  // create a trace to monitor the completion
+  const traceClient = new TraceClient();
+  const messageLength = messages.length;
+  const systemRole = messages.find((message) => message.role === 'system')?.content;
+
+  const trace = traceClient.createTrace({
+    id: tracePayload?.traceId,
+    input: messages,
+    metadata: { messageLength, model, provider, systemRole, tools },
+    name: tracePayload?.traceName,
+    sessionId: tracePayload?.topicId
+      ? tracePayload.topicId
+      : `${tracePayload?.sessionId || INBOX_SESSION_ID}@default`,
+    tags: tracePayload?.tags,
+    userId: tracePayload?.userId,
+  });
+
+  const generation = trace?.generation({
+    input: messages,
+    metadata: { messageLength, model, provider },
+    model,
+    modelParameters: parameters as any,
+    name: `Chat Completion (${provider})`,
+    startTime: new Date(),
+  });
+
+  return {
+    callback: {
+      onCompletion: async ({ text, thinking, usage, grounding, toolsCalling }) => {
+        const output =
+          // if the toolsCalling is not empty, we need to return the toolsCalling
+          !!toolsCalling && toolsCalling.length > 0
+            ? !!text
+              ? // tools calling with thinking and text
+                { text, thinking, toolsCalling }
+              : toolsCalling
+            : !!thinking
+              ? { text, thinking }
+              : text;
+
+        generation?.update({
+          endTime: new Date(),
+          metadata: { grounding, thinking },
+          output,
+          usage: usage
+            ? {
+                completionTokens: usage.outputTextTokens,
+                input: usage.totalInputTokens,
+                output: usage.totalOutputTokens,
+                promptTokens: usage.inputTextTokens,
+                totalTokens: usage.totalTokens,
+              }
+            : undefined,
+        });
+
+        trace?.update({ output });
+      },
+
+      onFinal: () => {
+        after(async () => {
+          try {
+            await traceClient.shutdownAsync();
+          } catch (e) {
+            console.error('TraceClient shutdown error:', e);
+          }
+        });
+      },
+
+      onStart: () => {
+        generation?.update({ completionStartTime: new Date() });
+      },
+
+      onToolsCalling: async () => {
+        trace?.update({
+          tags: [...(tracePayload?.tags || []), TraceTagMap.ToolsCalling],
+        });
+      },
+    } as ChatStreamCallbacks,
+    headers: {
+      [LOBE_CHAT_OBSERVATION_ID]: generation?.id,
+      [LOBE_CHAT_TRACE_ID]: trace?.id,
+    },
+  };
+};

package/src/store/chat/slices/aiChat/actions/generateAIChat.ts

@@ -388,6 +388,12 @@ export const generateAIChat: StateCreator<
             });
           }
         },
+        trace: {
+          traceId: params?.traceId,
+          sessionId: get().activeId,
+          topicId: get().activeTopicId,
+          traceName: TraceNameMap.SearchIntentRecognition,
+        },
         abortController,
         onMessageHandle: async (chunk) => {
           if (chunk.type === 'tool_calls') {

package/src/utils/fetch/fetchSSE.ts

@@ -2,6 +2,7 @@ import { isObject } from 'lodash-es';
 
 import { MESSAGE_CANCEL_FLAT } from '@/const/message';
 import { LOBE_CHAT_OBSERVATION_ID, LOBE_CHAT_TRACE_ID } from '@/const/trace';
+import { parseToolCalls } from '@/libs/agent-runtime';
 import { ChatErrorType } from '@/types/fetch';
 import { SmoothingParams } from '@/types/llm';
 import {
@@ -18,7 +19,6 @@ import { nanoid } from '@/utils/uuid';
 
 import { fetchEventSource } from './fetchEventSource';
 import { getMessageError } from './parseError';
-import { parseToolCalls } from './parseToolCalls';
 
 type SSEFinishType = 'done' | 'error' | 'abort';