@lobehub/chat 1.61.6 → 1.62.1

package/src/database/migrations/meta/_journal.json

@@ -105,6 +105,13 @@
       "when": 1737609172353,
       "tag": "0014_add_message_reasoning",
       "breakpoints": true
+    },
+    {
+      "idx": 15,
+      "version": "7",
+      "when": 1739901891891,
+      "tag": "0015_add_message_search_metadata",
+      "breakpoints": true
     }
   ],
   "version": "6"

package/src/database/schemas/message.ts

@@ -13,7 +13,7 @@ import {
 import { createSelectSchema } from 'drizzle-zod';
 
 import { idGenerator } from '@/database/utils/idGenerator';
-import { ModelReasoning } from '@/types/message';
+import { GroundingSearch, ModelReasoning } from '@/types/message';
 
 import { timestamps } from './_helpers';
 import { agents } from './agent';
@@ -34,6 +34,8 @@ export const messages = pgTable(
     role: text('role', { enum: ['user', 'system', 'assistant', 'tool'] }).notNull(),
     content: text('content'),
     reasoning: jsonb('reasoning').$type<ModelReasoning>(),
+    search: jsonb('search').$type<GroundingSearch>(),
+    metadata: jsonb('metadata'),
 
     model: text('model'),
     provider: text('provider'),

package/src/database/server/models/message.ts

@@ -74,6 +74,8 @@ export class MessageModel {
         role: messages.role,
         content: messages.content,
         reasoning: messages.reasoning,
+        search: messages.search,
+        metadata: messages.metadata,
         error: messages.error,
 
         model: messages.model,

package/src/features/Conversation/components/ChatItem/index.tsx

@@ -172,12 +172,21 @@ const Item = memo<ChatListItemProps>(
 
     const markdownProps = useMemo(
       () => ({
+        citations: item?.role === 'user' ? undefined : item?.search?.citations,
         components,
         customRender: markdownCustomRender,
         rehypePlugins: item?.role === 'user' ? undefined : rehypePlugins,
         remarkPlugins: item?.role === 'user' ? undefined : remarkPlugins,
+        showCitations:
+          item?.role === 'user'
+            ? undefined
+            : item?.search?.citations &&
+              // if the citations are all empty, we should not show the citations
+              item?.search?.citations.length > 0 &&
+              // if the citations's url and title are all the same, we should not show the citations
+              item?.search?.citations.every((item) => item.title !== item.url),
       }),
-      [components, markdownCustomRender, item?.role],
+      [components, markdownCustomRender, item?.role, item?.search],
     );
 
     const onChange = useCallback((value: string) => updateMessageContent(id, value), [id]);

package/src/features/Conversation/components/MarkdownElements/Thinking/Render.tsx

@@ -18,10 +18,14 @@ const Render = memo<MarkdownElementProps>(({ children, id }) => {
     const message = chatSelectors.getMessageById(id)(s);
     return [!isThinkingClosed(message?.content)];
   });
+  const citations = useChatStore((s) => {
+    const message = chatSelectors.getMessageById(id)(s);
+    return message?.search?.citations;
+  });
 
   if (!isGenerating && !children) return;
 
-  return <Thinking content={children as string} thinking={isGenerating} />;
+  return <Thinking citations={citations} content={children as string} thinking={isGenerating} />;
 });
 
 export default Render;

package/src/features/Conversation/components/MarkdownElements/remarkPlugins/createRemarkCustomTagPlugin.ts

@@ -32,6 +32,7 @@ export const createRemarkCustomTagPlugin = (tag: string) => () => {
         );
 
         // 转换为 Markdown 字符串
+
         const content = treeNodeToString(contentNodes);
 
         // 创建自定义节点

package/src/features/Conversation/components/MarkdownElements/remarkPlugins/getNodeContent.test.ts

@@ -393,4 +393,111 @@ describe('treeNodeToString', () => {
 1. 求 $a_2$ 和 $a_3$，根据前三项的规律猜想该数列的通项公式
 2. 用数学归纳法证明你的猜想。`);
   });
+
+  describe('link node', () => {
+    it('with url', () => {
+      const nodes = [
+        {
+          type: 'paragraph',
+          children: [
+            {
+              type: 'link',
+              title: null,
+              url: 'citation-1',
+              children: [
+                {
+                  type: 'text',
+                  value: '#citation-1',
+                  position: {
+                    start: {
+                      line: 5,
+                      column: 26,
+                      offset: 78,
+                    },
+                    end: {
+                      line: 5,
+                      column: 37,
+                      offset: 89,
+                    },
+                  },
+                },
+              ],
+              position: {
+                start: {
+                  line: 5,
+                  column: 25,
+                  offset: 77,
+                },
+                end: {
+                  line: 5,
+                  column: 50,
+                  offset: 102,
+                },
+              },
+            },
+          ],
+          position: {
+            start: {
+              line: 5,
+              column: 1,
+              offset: 53,
+            },
+            end: {
+              line: 5,
+              column: 220,
+              offset: 272,
+            },
+          },
+        },
+      ];
+
+      const result = treeNodeToString(nodes as Parent[]);
+
+      expect(result).toEqual(`[#citation-1](citation-1)`);
+    });
+
+    it('handle error case', () => {
+      const nodes = [
+        {
+          type: 'paragraph',
+          children: [
+            {
+              type: 'link',
+              title: null,
+              url: 'citation-1',
+              children: [],
+              position: {
+                start: {
+                  line: 5,
+                  column: 25,
+                  offset: 77,
+                },
+                end: {
+                  line: 5,
+                  column: 50,
+                  offset: 102,
+                },
+              },
+            },
+          ],
+          position: {
+            start: {
+              line: 5,
+              column: 1,
+              offset: 53,
+            },
+            end: {
+              line: 5,
+              column: 220,
+              offset: 272,
+            },
+          },
+        },
+      ];
+
+      const result = treeNodeToString(nodes as Parent[]);
+
+      expect(result).toEqual(`[](citation-1)`);
+    });
+  });
 });

package/src/features/Conversation/components/MarkdownElements/remarkPlugins/getNodeContent.ts

@@ -7,6 +7,12 @@ const processNode = (node: any): string => {
     return `$${node.value}$`;
   }
 
+  if (node.type === 'link') {
+    const text = node.children?.[0] ? processNode(node.children?.[0]) : '';
+
+    return `[${text}](${node.url})`;
+  }
+
   // 处理带有子节点的容器
   if (node.children) {
     const content = node.children.map((element: Parent) => processNode(element)).join('');

package/src/libs/agent-runtime/perplexity/index.test.ts

@@ -28,17 +28,9 @@ beforeEach(() => {
   );
 });
 
-afterEach(() => {
-  vi.clearAllMocks();
-});
-
 describe('LobePerplexityAI', () => {
   describe('chat', () => {
     it('should call chat method with temperature', async () => {
-      vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
-        new ReadableStream() as any,
-      );
-
       await instance.chat({
         messages: [{ content: 'Hello', role: 'user' }],
         model: 'text-davinci-003',
@@ -56,10 +48,6 @@ describe('LobePerplexityAI', () => {
     });
 
     it('should be undefined when temperature >= 2', async () => {
-      vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
-        new ReadableStream() as any,
-      );
-
       await instance.chat({
         messages: [{ content: 'Hello', role: 'user' }],
         model: 'text-davinci-003',
@@ -75,5 +63,161 @@ describe('LobePerplexityAI', () => {
         expect.any(Object),
       );
     });
+
+    it('should with search citations', async () => {
+      const data = [
+        {
+          id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
+          model: 'sonar-pro',
+          created: 1739896615,
+          usage: {
+            prompt_tokens: 4,
+            completion_tokens: 3,
+            total_tokens: 7,
+            citation_tokens: 2217,
+            num_search_queries: 1,
+          },
+          citations: [
+            'https://www.weather.com.cn/weather/101210101.shtml',
+            'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
+            'https://weather.cma.cn/web/weather/58457.html',
+            'https://tianqi.so.com/weather/101210101',
+            'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
+            'https://www.hzqx.com',
+            'https://www.hzqx.com/pc/hztq/',
+          ],
+          object: 'chat.completion',
+          choices: [
+            {
+              index: 0,
+              finish_reason: null,
+              message: {
+                role: 'assistant',
+                content: '杭州今',
+              },
+              delta: {
+                role: 'assist',
+                content: '杭州今',
+              },
+            },
+          ],
+        },
+        {
+          id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
+          model: 'sonar-pro',
+          created: 1739896615,
+          usage: {
+            prompt_tokens: 4,
+            completion_tokens: 9,
+            total_tokens: 13,
+            citation_tokens: 2217,
+            num_search_queries: 1,
+          },
+          citations: [
+            'https://www.weather.com.cn/weather/101210101.shtml',
+            'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
+            'https://weather.cma.cn/web/weather/58457.html',
+            'https://tianqi.so.com/weather/101210101',
+            'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
+            'https://www.hzqx.com',
+            'https://www.hzqx.com/pc/hztq/',
+          ],
+          object: 'chat.completion',
+          choices: [
+            {
+              index: 0,
+              finish_reason: null,
+              message: {
+                role: 'assistant',
+                content: '杭州今天和未来几天的',
+              },
+              delta: {
+                role: 'assistant',
+                content: '天和未来几天的',
+              },
+            },
+          ],
+        },
+        {
+          id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
+          model: 'sonar-pro',
+          created: 1739896615,
+          usage: {
+            prompt_tokens: 4,
+            completion_tokens: 14,
+            total_tokens: 18,
+            citation_tokens: 2217,
+            num_search_queries: 1,
+          },
+          citations: [
+            'https://www.weather.com.cn/weather/101210101.shtml',
+            'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
+            'https://weather.cma.cn/web/weather/58457.html',
+            'https://tianqi.so.com/weather/101210101',
+            'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
+            'https://www.hzqx.com',
+            'https://www.hzqx.com/pc/hztq/',
+          ],
+          object: 'chat.completion',
+          choices: [
+            {
+              index: 0,
+              finish_reason: null,
+              message: {
+                role: 'assistant',
+                content: '杭州今天和未来几天的天气预报如',
+              },
+            },
+          ],
+        },
+      ];
+
+      const mockStream = new ReadableStream({
+        start(controller) {
+          data.forEach((chunk) => {
+            controller.enqueue(chunk);
+          });
+
+          controller.close();
+        },
+      });
+
+      vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(mockStream as any);
+
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0,
+      });
+
+      const decoder = new TextDecoder();
+      const reader = result.body!.getReader();
+      const stream: string[] = [];
+
+      while (true) {
+        const { value, done } = await reader.read();
+        if (done) break;
+        stream.push(decoder.decode(value));
+      }
+
+      expect(stream).toEqual(
+        [
+          'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
+          'event: citations',
+          'data: [{"title":"https://www.weather.com.cn/weather/101210101.shtml","url":"https://www.weather.com.cn/weather/101210101.shtml"},{"title":"https://tianqi.moji.com/weather/china/zhejiang/hangzhou","url":"https://tianqi.moji.com/weather/china/zhejiang/hangzhou"},{"title":"https://weather.cma.cn/web/weather/58457.html","url":"https://weather.cma.cn/web/weather/58457.html"},{"title":"https://tianqi.so.com/weather/101210101","url":"https://tianqi.so.com/weather/101210101"},{"title":"https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832","url":"https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832"},{"title":"https://www.hzqx.com","url":"https://www.hzqx.com"},{"title":"https://www.hzqx.com/pc/hztq/","url":"https://www.hzqx.com/pc/hztq/"}]\n',
+          'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
+          'event: text',
+          'data: "杭州今"\n',
+          'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
+          'event: text',
+          'data: "天和未来几天的"\n',
+          'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
+          'event: data',
+          'data: {"id":"506d64fb-e7f2-4d94-b80f-158369e9446d","index":0}\n',
+        ].map((line) => `${line}\n`),
+      );
+
+      expect((await reader.read()).done).toBe(true);
+    });
   });
 });

package/src/libs/agent-runtime/utils/streams/anthropic.ts

@@ -3,9 +3,9 @@ import type { Stream } from '@anthropic-ai/sdk/streaming';
 
 import { ChatStreamCallbacks } from '../../types';
 import {
+  StreamContext,
   StreamProtocolChunk,
   StreamProtocolToolCallChunk,
-  StreamStack,
   StreamToolCallChunkData,
   convertIterableToStream,
   createCallbacksTransformer,
@@ -14,7 +14,7 @@ import {
 
 export const transformAnthropicStream = (
   chunk: Anthropic.MessageStreamEvent,
-  stack: StreamStack,
+  stack: StreamContext,
 ): StreamProtocolChunk => {
   // maybe need another structure to add support for multiple choices
   switch (chunk.type) {
@@ -100,7 +100,7 @@ export const AnthropicStream = (
   stream: Stream<Anthropic.MessageStreamEvent> | ReadableStream,
   callbacks?: ChatStreamCallbacks,
 ) => {
-  const streamStack: StreamStack = { id: '' };
+  const streamStack: StreamContext = { id: '' };
 
   const readableStream =
     stream instanceof ReadableStream ? stream : convertIterableToStream(stream);

package/src/libs/agent-runtime/utils/streams/bedrock/claude.ts

@@ -4,14 +4,18 @@ import { nanoid } from '@/utils/uuid';
 
 import { ChatStreamCallbacks } from '../../../types';
 import { transformAnthropicStream } from '../anthropic';
-import { StreamStack, createCallbacksTransformer, createSSEProtocolTransformer } from '../protocol';
+import {
+  StreamContext,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+} from '../protocol';
 import { createBedrockStream } from './common';
 
 export const AWSBedrockClaudeStream = (
   res: InvokeModelWithResponseStreamResponse | ReadableStream,
   cb?: ChatStreamCallbacks,
 ): ReadableStream<string> => {
-  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+  const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
 

package/src/libs/agent-runtime/utils/streams/bedrock/llama.ts

@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
 
 import { ChatStreamCallbacks } from '../../../types';
 import {
+  StreamContext,
   StreamProtocolChunk,
-  StreamStack,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
 } from '../protocol';
@@ -27,7 +27,7 @@ interface BedrockLlamaStreamChunk {
 
 export const transformLlamaStream = (
   chunk: BedrockLlamaStreamChunk,
-  stack: StreamStack,
+  stack: StreamContext,
 ): StreamProtocolChunk => {
   // maybe need another structure to add support for multiple choices
   if (chunk.stop_reason) {
@@ -41,7 +41,7 @@ export const AWSBedrockLlamaStream = (
   res: InvokeModelWithResponseStreamResponse | ReadableStream,
   cb?: ChatStreamCallbacks,
 ): ReadableStream<string> => {
-  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+  const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
 

package/src/libs/agent-runtime/utils/streams/google-ai.ts

@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
 
 import { ChatStreamCallbacks } from '../../types';
 import {
+  StreamContext,
   StreamProtocolChunk,
-  StreamStack,
   StreamToolCallChunkData,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
@@ -14,7 +14,7 @@ import {
 
 const transformGoogleGenerativeAIStream = (
   chunk: EnhancedGenerateContentResponse,
-  stack: StreamStack,
+  stack: StreamContext,
 ): StreamProtocolChunk => {
   // maybe need another structure to add support for multiple choices
   const functionCalls = chunk.functionCalls();
@@ -49,7 +49,7 @@ export const GoogleGenerativeAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
   callbacks?: ChatStreamCallbacks,
 ) => {
-  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+  const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
     .pipeThrough(createSSEProtocolTransformer(transformGoogleGenerativeAIStream, streamStack))

package/src/libs/agent-runtime/utils/streams/ollama.ts

@@ -4,14 +4,14 @@ import { ChatStreamCallbacks } from '@/libs/agent-runtime';
 import { nanoid } from '@/utils/uuid';
 
 import {
+  StreamContext,
   StreamProtocolChunk,
-  StreamStack,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
   generateToolCallId,
 } from './protocol';
 
-const transformOllamaStream = (chunk: ChatResponse, stack: StreamStack): StreamProtocolChunk => {
+const transformOllamaStream = (chunk: ChatResponse, stack: StreamContext): StreamProtocolChunk => {
   // maybe need another structure to add support for multiple choices
   if (chunk.done && !chunk.message.content) {
     return { data: 'finished', id: stack.id, type: 'stop' };
@@ -39,7 +39,7 @@ export const OllamaStream = (
   res: ReadableStream<ChatResponse>,
   cb?: ChatStreamCallbacks,
 ): ReadableStream<string> => {
-  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+  const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return res
     .pipeThrough(createSSEProtocolTransformer(transformOllamaStream, streamStack))

package/src/libs/agent-runtime/utils/streams/openai.ts

@@ -1,15 +1,15 @@
 import OpenAI from 'openai';
 import type { Stream } from 'openai/streaming';
 
-import { ChatMessageError } from '@/types/message';
+import { ChatMessageError, CitationItem } from '@/types/message';
 
 import { AgentRuntimeErrorType, ILobeAgentRuntimeErrorType } from '../../error';
 import { ChatStreamCallbacks } from '../../types';
 import {
   FIRST_CHUNK_ERROR_KEY,
+  StreamContext,
   StreamProtocolChunk,
   StreamProtocolToolCallChunk,
-  StreamStack,
   StreamToolCallChunkData,
   convertIterableToStream,
   createCallbacksTransformer,
@@ -20,8 +20,8 @@ import {
 
 export const transformOpenAIStream = (
   chunk: OpenAI.ChatCompletionChunk,
-  stack?: StreamStack,
-): StreamProtocolChunk => {
+  streamContext: StreamContext,
+): StreamProtocolChunk | StreamProtocolChunk[] => {
   // handle the first chunk error
   if (FIRST_CHUNK_ERROR_KEY in chunk) {
     delete chunk[FIRST_CHUNK_ERROR_KEY];
@@ -48,8 +48,8 @@ export const transformOpenAIStream = (
     if (typeof item.delta?.tool_calls === 'object' && item.delta.tool_calls?.length > 0) {
       return {
         data: item.delta.tool_calls.map((value, index): StreamToolCallChunkData => {
-          if (stack && !stack.tool) {
-            stack.tool = { id: value.id!, index: value.index, name: value.function!.name! };
+          if (streamContext && !streamContext.tool) {
+            streamContext.tool = { id: value.id!, index: value.index, name: value.function!.name! };
           }
 
           return {
@@ -57,7 +57,10 @@ export const transformOpenAIStream = (
               arguments: value.function?.arguments ?? '{}',
               name: value.function?.name ?? null,
             },
-            id: value.id || stack?.tool?.id || generateToolCallId(index, value.function?.name),
+            id:
+              value.id ||
+              streamContext?.tool?.id ||
+              generateToolCallId(index, value.function?.name),
 
             // mistral's tool calling don't have index and function field, it's data like:
             // [{"id":"xbhnmTtY7","function":{"name":"lobe-image-designer____text2image____builtin","arguments":"{\"prompts\": [\"A photo of a small, fluffy dog with a playful expression and wagging tail.\", \"A watercolor painting of a small, energetic dog with a glossy coat and bright eyes.\", \"A vector illustration of a small, adorable dog with a short snout and perky ears.\", \"A drawing of a small, scruffy dog with a mischievous grin and a wagging tail.\"], \"quality\": \"standard\", \"seeds\": [123456, 654321, 111222, 333444], \"size\": \"1024x1024\", \"style\": \"vivid\"}"}}]
@@ -114,6 +117,21 @@ export const transformOpenAIStream = (
       }
 
       if (typeof content === 'string') {
+        // in Perplexity api, the citation is in every chunk, but we only need to return it once
+
+        if ('citations' in chunk && !streamContext?.returnedPplxCitation) {
+          streamContext.returnedPplxCitation = true;
+
+          const citations = (chunk.citations as any[]).map((item) =>
+            typeof item === 'string' ? ({ title: item, url: item } as CitationItem) : item,
+          );
+
+          return [
+            { data: citations, id: chunk.id, type: 'citations' },
+            { data: content, id: chunk.id, type: 'text' },
+          ];
+        }
+
         return { data: content, id: chunk.id, type: 'text' };
       }
     }
@@ -164,7 +182,7 @@ export const OpenAIStream = (
   stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
   { callbacks, provider, bizErrorTypeTransformer }: OpenAIStreamOptions = {},
 ) => {
-  const streamStack: StreamStack = { id: '' };
+  const streamStack: StreamContext = { id: '' };
 
   const readableStream =
     stream instanceof ReadableStream ? stream : convertIterableToStream(stream);