diff --git a/next.config.ts b/next.config.ts index 1cd1d965f1..ad5a22e678 100644 --- a/next.config.ts +++ b/next.config.ts @@ -42,7 +42,6 @@ const nextConfig: NextConfig = { // so we need to disable it // refs: https://github.com/lobehub/lobe-chat/pull/7430 serverMinification: false, - turbopackFileSystemCacheForDev: true, webVitalsAttribution: ['CLS', 'LCP'], webpackBuildWorker: true, webpackMemoryOptimizations: true, diff --git a/package.json b/package.json index d5bbc2d34b..04f69f4b6d 100644 --- a/package.json +++ b/package.json @@ -141,7 +141,7 @@ "@emotion/react": "^11.14.0", "@fal-ai/client": "^1.7.2", "@formkit/auto-animate": "^0.9.0", - "@google/genai": "^1.29.1", + "@google/genai": "^1.30.0", "@huggingface/inference": "^4.13.3", "@icons-pack/react-simple-icons": "^13.8.0", "@khmyznikov/pwa-install": "0.3.9", diff --git a/packages/context-engine/src/processors/ToolCall.ts b/packages/context-engine/src/processors/ToolCall.ts index 831f2da2b5..54d266e24f 100644 --- a/packages/context-engine/src/processors/ToolCall.ts +++ b/packages/context-engine/src/processors/ToolCall.ts @@ -130,6 +130,7 @@ export class ToolCallProcessor extends BaseProcessor { : `${tool.identifier}.${tool.apiName}`, }, id: tool.id, + thoughtSignature: tool.thoughtSignature, type: 'function', }), ); diff --git a/packages/context-engine/src/processors/__tests__/ToolCall.test.ts b/packages/context-engine/src/processors/__tests__/ToolCall.test.ts index dc78c4cf2c..cdf87ac51c 100644 --- a/packages/context-engine/src/processors/__tests__/ToolCall.test.ts +++ b/packages/context-engine/src/processors/__tests__/ToolCall.test.ts @@ -72,6 +72,65 @@ describe('ToolCallProcessor', () => { ]); }); + it('should pass through thoughtSignature when present', async () => { + const processor = new ToolCallProcessor(defaultConfig); + const context = createContext([ + { + content: '', + id: 'msg1', + role: 'assistant', + tools: [ + { + apiName: 'search', + arguments: '{"query":"test"}', + id: 'call_1', + identifier: 'web', + thoughtSignature: 'Let me search for this information', + type: 'builtin', + }, + ], + }, + ]); + + const result = await processor.process(context); + + expect(result.messages[0].tool_calls).toEqual([ + { + function: { + arguments: '{"query":"test"}', + name: 'web.search', + }, + id: 'call_1', + thoughtSignature: 'Let me search for this information', + type: 'function', + }, + ]); + }); + + it('should handle missing thoughtSignature', async () => { + const processor = new ToolCallProcessor(defaultConfig); + const context = createContext([ + { + content: '', + id: 'msg1', + role: 'assistant', + tools: [ + { + apiName: 'search', + arguments: '{"query":"test"}', + id: 'call_1', + identifier: 'web', + type: 'builtin', + }, + ], + }, + ]); + + const result = await processor.process(context); + + expect(result.messages[0].tool_calls[0].thoughtSignature).toBeUndefined(); + }); + it('should use custom genToolCallingName function', async () => { const genToolCallingName = vi.fn( (identifier, apiName, type) => `custom_${identifier}_${apiName}_${type}`, diff --git a/packages/context-engine/src/tools/ToolNameResolver.ts b/packages/context-engine/src/tools/ToolNameResolver.ts index 5bd3fb2b1d..684a017800 100644 --- a/packages/context-engine/src/tools/ToolNameResolver.ts +++ b/packages/context-engine/src/tools/ToolNameResolver.ts @@ -82,6 +82,7 @@ export class ToolNameResolver { arguments: toolCall.function.arguments, id: toolCall.id, identifier, + thoughtSignature: toolCall.thoughtSignature, type: (type ?? 'default') as any, }; diff --git a/packages/context-engine/src/tools/__tests__/ToolNameResolver.test.ts b/packages/context-engine/src/tools/__tests__/ToolNameResolver.test.ts index 6c2132e037..e4615d37e0 100644 --- a/packages/context-engine/src/tools/__tests__/ToolNameResolver.test.ts +++ b/packages/context-engine/src/tools/__tests__/ToolNameResolver.test.ts @@ -455,6 +455,63 @@ describe('ToolNameResolver', () => { }); }); + describe('resolve - thoughtSignature', () => { + it('should pass through thoughtSignature when present', () => { + const toolCalls = [ + { + function: { + arguments: '{"query": "test"}', + name: 'test-plugin____myAction____builtin', + }, + id: 'call_1', + thoughtSignature: 'thinking about this...', + type: 'function', + }, + ]; + + const manifests = { + 'test-plugin': { + api: [{ description: 'My action', name: 'myAction', parameters: {} }], + identifier: 'test-plugin', + meta: {}, + type: 'builtin' as const, + }, + }; + + const result = resolver.resolve(toolCalls, manifests); + + expect(result).toHaveLength(1); + expect(result[0].thoughtSignature).toBe('thinking about this...'); + }); + + it('should handle missing thoughtSignature', () => { + const toolCalls = [ + { + function: { + arguments: '{"query": "test"}', + name: 'test-plugin____myAction____builtin', + }, + id: 'call_1', + type: 'function', + }, + ]; + + const manifests = { + 'test-plugin': { + api: [{ description: 'My action', name: 'myAction', parameters: {} }], + identifier: 'test-plugin', + meta: {}, + type: 'builtin' as const, + }, + }; + + const result = resolver.resolve(toolCalls, manifests); + + expect(result).toHaveLength(1); + expect(result[0].thoughtSignature).toBeUndefined(); + }); + }); + describe('resolve - edge cases', () => { it('should filter out invalid tool calls with missing apiName', () => { const toolCalls = [ diff --git a/packages/context-engine/src/types.ts b/packages/context-engine/src/types.ts index 818d940202..5c43a525da 100644 --- a/packages/context-engine/src/types.ts +++ b/packages/context-engine/src/types.ts @@ -30,6 +30,7 @@ export interface MessageToolCall { name: string; }; id: string; + thoughtSignature?: string; type: 'function'; } export interface Message { diff --git a/packages/fetch-sse/src/fetchSSE.ts b/packages/fetch-sse/src/fetchSSE.ts index a3e3bf50c0..fc8f8a837a 100644 --- a/packages/fetch-sse/src/fetchSSE.ts +++ b/packages/fetch-sse/src/fetchSSE.ts @@ -17,7 +17,7 @@ import { nanoid } from '@lobechat/utils/uuid'; import { getMessageError } from './parseError'; -type SSEFinishType = 'done' | 'error' | 'abort'; +type SSEFinishType = 'done' | 'error' | 'abort' | string; export type OnFinishHandler = ( text: string, @@ -48,6 +48,10 @@ export interface MessageTextChunk { text: string; type: 'text'; } +export interface MessageStopChunk { + reason: string; + type: 'stop'; +} export interface MessageBase64ImageChunk { id: string; @@ -86,7 +90,8 @@ export interface FetchSSEOptions { | MessageGroundingChunk | MessageUsageChunk | MessageBase64ImageChunk - | MessageSpeedChunk, + | MessageSpeedChunk + | MessageStopChunk, ) => void; responseAnimation?: ResponseAnimation; } @@ -387,6 +392,11 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio break; } + case 'stop': { + options.onMessageHandle?.({ reason: data, type: 'stop' }); + break; + } + case 'reasoning': { if (textSmoothing) { thinkingController.pushToQueue(data); diff --git a/packages/model-runtime/src/core/contextBuilders/google.test.ts b/packages/model-runtime/src/core/contextBuilders/google.test.ts index 145d800990..2e9e1fa4b2 100644 --- a/packages/model-runtime/src/core/contextBuilders/google.test.ts +++ b/packages/model-runtime/src/core/contextBuilders/google.test.ts @@ -6,6 +6,7 @@ import { describe, expect, it, vi } from 'vitest'; import { ChatCompletionTool, OpenAIChatMessage, UserMessageContentPart } from '../../types'; import { parseDataUri } from '../../utils/uriParser'; import { + GEMINI_MAGIC_THOUGHT_SIGNATURE, buildGoogleMessage, buildGoogleMessages, buildGooglePart, @@ -232,6 +233,415 @@ describe('google contextBuilders', () => { }); }); + it('should correctly convert function call message with thoughtSignature', async () => { + const message = { + role: 'assistant', + tool_calls: [ + { + function: { + arguments: JSON.stringify({ + language: ['JSON'], + path: 'package.json', + query: '"version":', + repo: 'lobehub/lobe-chat', + }), + name: 'grep____searchGitHub____mcp', + }, + id: 'grep____searchGitHub____mcp_0_6RnOMTF0', + thoughtSignature: + 'EsUHCsIHAdHtim9/MrjP+pnhM8DVkvulyfWQVf+isXQxEAbF32gbflE1hl6Te80qtp77Ywn8opB2uhQOIH/l6SStsj3+XRy1U1DTeKtqZxDBoLP2rNK6pi3/nk0ZOQIc8f6rxB70G/zOhk7d/1XQFqhmw5H+yDVRQjGD1cNPY5ctWGxQLAIk/HMWNovUJzz2c81jGWoXu7k2vtpuur2hcAL+J79BEVUTfvU3mSiXqJFTClmFPB6Fe79i0y3TwM2XdIBxzPgVgf8B+Pnv1S6YDxHNSm46jTlXKcSw30r3ixs5xEOzerbOUW5WG9BGukw/YQVvHiuoGLIALRa2Ig7dlOMH8+o+f0mKJtyYj8yF6wyBMol+G4mhSHvQSKJLj/Z5kFHvDZKeVUEOZed6vZivYLrVezjQPXgLHJMOmbp6QrZGxqW45QxDKY5X5F8giIOM8VgsUYhDQUBown+3vvwkIBA24icDsOwdhJ/roe9GabbGfxpkSzARIFh7rSI01cRKbh6cEaVFXf2WQftPeD7dBseQLiCdUYoy4ytECrjTpknrWnVUG6Ly4SKW6uN/IJXpm9JT9GgnGLIddFtEQzm9sIKWNpGEz6++lZpiCFS6LsYSnTP3vPj/7oSABRmwWywxA8EmLh+sv+jiK5aMjFi1sTuJ0Ujsvza3/SHZKewNi9WKQUDOa9Mqtjs2YGDnJxto4l5GMUzI5vhf6/+/A5eHALfVabaFP97v8FEPrXQU94dognwx4EnNqy/KWmGIlYZYqIfjaSAy7Z74viwl+oTtL9gyyBDc/FrQvXfyrYIq8N0pkLKAEh33fa/+YVocLL1LKI9rb2bg/RRr+Ee4NyIQKhIdEJaEh74d1COd/4r06J92ThkfVo5PEVTSsr8tBKiJ5wSmX9vyhbLWzxmXoq1xfGrs8kg7NMW53XEWGlQrIVOQmUtjjjBQKj6b4rBTAO6EKk63cGFbkSPohifiUBPHbxUUPy/hf0tQpeOo3jA01AuCFLOIZ5IYJ+Rm5+aZTU3Panv+Q7Yl1w5t5swhbNZfg7MlU/sxwLijLuWDDNfw+2Zw/aa3VDPgVw6Nv2vKkHi4tUU0XlgfiQgQYUMPxpGRV837uUxvZFNep2QUlAMog5h4sMYJWIAX1kK1pzsyR/KxuCn6nUq4ovWNBQHLC4aW2ZcGgW/6CbF81F1cewUz+vWNMMkJrL0d9celGEbFuY0Q709UipaDbCg49twlnLV9XUwqC5wYTFBiJbynBDqiZAvXn2YOxNIs8CCzuu2GSCQDo09ksJy5g/o=', + type: 'function', + }, + ], + } as OpenAIChatMessage; + + const converted = await buildGoogleMessage(message); + + expect(converted).toEqual({ + parts: [ + { + functionCall: { + args: { + language: ['JSON'], + path: 'package.json', + query: '"version":', + repo: 'lobehub/lobe-chat', + }, + name: 'grep____searchGitHub____mcp', + }, + thoughtSignature: + 'EsUHCsIHAdHtim9/MrjP+pnhM8DVkvulyfWQVf+isXQxEAbF32gbflE1hl6Te80qtp77Ywn8opB2uhQOIH/l6SStsj3+XRy1U1DTeKtqZxDBoLP2rNK6pi3/nk0ZOQIc8f6rxB70G/zOhk7d/1XQFqhmw5H+yDVRQjGD1cNPY5ctWGxQLAIk/HMWNovUJzz2c81jGWoXu7k2vtpuur2hcAL+J79BEVUTfvU3mSiXqJFTClmFPB6Fe79i0y3TwM2XdIBxzPgVgf8B+Pnv1S6YDxHNSm46jTlXKcSw30r3ixs5xEOzerbOUW5WG9BGukw/YQVvHiuoGLIALRa2Ig7dlOMH8+o+f0mKJtyYj8yF6wyBMol+G4mhSHvQSKJLj/Z5kFHvDZKeVUEOZed6vZivYLrVezjQPXgLHJMOmbp6QrZGxqW45QxDKY5X5F8giIOM8VgsUYhDQUBown+3vvwkIBA24icDsOwdhJ/roe9GabbGfxpkSzARIFh7rSI01cRKbh6cEaVFXf2WQftPeD7dBseQLiCdUYoy4ytECrjTpknrWnVUG6Ly4SKW6uN/IJXpm9JT9GgnGLIddFtEQzm9sIKWNpGEz6++lZpiCFS6LsYSnTP3vPj/7oSABRmwWywxA8EmLh+sv+jiK5aMjFi1sTuJ0Ujsvza3/SHZKewNi9WKQUDOa9Mqtjs2YGDnJxto4l5GMUzI5vhf6/+/A5eHALfVabaFP97v8FEPrXQU94dognwx4EnNqy/KWmGIlYZYqIfjaSAy7Z74viwl+oTtL9gyyBDc/FrQvXfyrYIq8N0pkLKAEh33fa/+YVocLL1LKI9rb2bg/RRr+Ee4NyIQKhIdEJaEh74d1COd/4r06J92ThkfVo5PEVTSsr8tBKiJ5wSmX9vyhbLWzxmXoq1xfGrs8kg7NMW53XEWGlQrIVOQmUtjjjBQKj6b4rBTAO6EKk63cGFbkSPohifiUBPHbxUUPy/hf0tQpeOo3jA01AuCFLOIZ5IYJ+Rm5+aZTU3Panv+Q7Yl1w5t5swhbNZfg7MlU/sxwLijLuWDDNfw+2Zw/aa3VDPgVw6Nv2vKkHi4tUU0XlgfiQgQYUMPxpGRV837uUxvZFNep2QUlAMog5h4sMYJWIAX1kK1pzsyR/KxuCn6nUq4ovWNBQHLC4aW2ZcGgW/6CbF81F1cewUz+vWNMMkJrL0d9celGEbFuY0Q709UipaDbCg49twlnLV9XUwqC5wYTFBiJbynBDqiZAvXn2YOxNIs8CCzuu2GSCQDo09ksJy5g/o=', + }, + ], + role: 'model', + }); + }); + + describe('should correctly convert function call message without thoughtSignature', () => { + it('should add magic signature when last message is tool message', async () => { + const messages: OpenAIChatMessage[] = [ + { + content: 'Web Browsing plugin available', + role: 'system', + }, + { + content: '杭州天气如何', + role: 'user', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"杭州天气","searchEngines":["google"]}', + name: 'lobe-web-browsing____search____builtin', + }, + id: 'call_001', + type: 'function', + }, + ], + }, + { + content: 'Tool execution was aborted by user.', + name: 'lobe-web-browsing____search____builtin', + role: 'tool', + tool_call_id: 'call_001', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"杭州 天气","searchEngines":["bing"]}', + name: 'lobe-web-browsing____search____builtin', + }, + id: 'call_002', + type: 'function', + }, + ], + }, + { + content: 'no result', + name: 'lobe-web-browsing____search____builtin', + role: 'tool', + tool_call_id: 'call_002', + }, + ]; + + const contents = await buildGoogleMessages(messages); + + expect(contents).toEqual([ + { + parts: [{ text: 'Web Browsing plugin available' }], + role: 'user', + }, + { parts: [{ text: '杭州天气如何' }], role: 'user' }, + { + parts: [ + { + functionCall: { + args: { query: '杭州天气', searchEngines: ['google'] }, + name: 'lobe-web-browsing____search____builtin', + }, + thoughtSignature: GEMINI_MAGIC_THOUGHT_SIGNATURE, + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'lobe-web-browsing____search____builtin', + response: { result: 'Tool execution was aborted by user.' }, + }, + }, + ], + role: 'user', + }, + { + parts: [ + { + functionCall: { + args: { query: '杭州 天气', searchEngines: ['bing'] }, + name: 'lobe-web-browsing____search____builtin', + }, + thoughtSignature: GEMINI_MAGIC_THOUGHT_SIGNATURE, + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'lobe-web-browsing____search____builtin', + response: { result: 'no result' }, + }, + }, + ], + role: 'user', + }, + ]); + }); + + it('should NOT add magic signature when thoughtSignature already exists', async () => { + const existingSignature = 'existing_signature_from_model'; + const messages: OpenAIChatMessage[] = [ + { + content: '杭州天气如何', + role: 'user', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"杭州天气","searchEngines":["google"]}', + name: 'lobe-web-browsing____search____builtin', + }, + id: 'call_001', + thoughtSignature: existingSignature, + type: 'function', + }, + ], + }, + { + content: 'Tool result', + name: 'lobe-web-browsing____search____builtin', + role: 'tool', + tool_call_id: 'call_001', + }, + ]; + + const contents = await buildGoogleMessages(messages); + + expect(contents).toEqual([ + { + parts: [{ text: '杭州天气如何' }], + role: 'user', + }, + { + parts: [ + { + functionCall: { + args: { query: '杭州天气', searchEngines: ['google'] }, + name: 'lobe-web-browsing____search____builtin', + }, + // Should keep existing thoughtSignature, not add magic signature + thoughtSignature: existingSignature, + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'lobe-web-browsing____search____builtin', + response: { result: 'Tool result' }, + }, + }, + ], + role: 'user', + }, + ]); + }); + + it('should add magic signature only after last user message in multi-turn scenario', async () => { + const messages: OpenAIChatMessage[] = [ + { + content: 'First question', + role: 'user', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"first"}', + name: 'search', + }, + id: 'call_001', + type: 'function', + }, + ], + }, + { + content: 'First result', + name: 'search', + role: 'tool', + tool_call_id: 'call_001', + }, + { + content: 'Second question', + role: 'user', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"second"}', + name: 'search', + }, + id: 'call_002', + type: 'function', + }, + ], + }, + { + content: 'Second result', + name: 'search', + role: 'tool', + tool_call_id: 'call_002', + }, + ]; + + const contents = await buildGoogleMessages(messages); + + expect(contents).toEqual([ + { + parts: [{ text: 'First question' }], + role: 'user', + }, + { + parts: [ + { + functionCall: { + args: { query: 'first' }, + name: 'search', + }, + // No magic signature for this one (before last user message) + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'search', + response: { result: 'First result' }, + }, + }, + ], + role: 'user', + }, + { + parts: [{ text: 'Second question' }], + role: 'user', + }, + { + parts: [ + { + functionCall: { + args: { query: 'second' }, + name: 'search', + }, + // Magic signature added (after last user message) + thoughtSignature: GEMINI_MAGIC_THOUGHT_SIGNATURE, + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'search', + response: { result: 'Second result' }, + }, + }, + ], + role: 'user', + }, + ]); + }); + + it('should NOT add magic signature when last message is user text message', async () => { + const messages: OpenAIChatMessage[] = [ + { + content: 'Web Browsing plugin available', + role: 'system', + }, + { + content: '杭州天气如何', + role: 'user', + }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: '{"query":"杭州天气","searchEngines":["google"]}', + name: 'lobe-web-browsing____search____builtin', + }, + id: 'call_001', + type: 'function', + }, + ], + }, + { + content: 'Tool execution was aborted by user.', + name: 'lobe-web-browsing____search____builtin', + role: 'tool', + tool_call_id: 'call_001', + }, + { + content: 'Please try again', + role: 'user', + }, + ]; + + const contents = await buildGoogleMessages(messages); + + expect(contents).toEqual([ + { + parts: [{ text: 'Web Browsing plugin available' }], + role: 'user', + }, + { + parts: [{ text: '杭州天气如何' }], + role: 'user', + }, + { + parts: [ + { + functionCall: { + args: { query: '杭州天气', searchEngines: ['google'] }, + name: 'lobe-web-browsing____search____builtin', + }, + // No thoughtSignature should be added when last message is user text + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'lobe-web-browsing____search____builtin', + response: { result: 'Tool execution was aborted by user.' }, + }, + }, + ], + role: 'user', + }, + { + parts: [{ text: 'Please try again' }], + role: 'user', + }, + ]); + }); + }); + it('should correctly handle empty content', async () => { const message: OpenAIChatMessage = { content: '' as any, // explicitly set as empty string @@ -361,6 +771,7 @@ describe('google contextBuilders', () => { args: { location: 'London', unit: 'celsius' }, name: 'get_current_weather', }, + thoughtSignature: GEMINI_MAGIC_THOUGHT_SIGNATURE, }, ], role: 'model', @@ -410,6 +821,74 @@ describe('google contextBuilders', () => { { parts: [{ text: 'Hi' }], role: 'model' }, ]); }); + + it('should correctly convert full conversation with thoughtSignature', async () => { + const messages: OpenAIChatMessage[] = [ + { content: 'system prompt', role: 'system' }, + { content: 'LobeChat 最新版本', role: 'user' }, + { + content: '', + role: 'assistant', + tool_calls: [ + { + function: { + arguments: JSON.stringify({ + language: ['JSON'], + path: 'package.json', + query: '"version":', + repo: 'lobehub/lobe-chat', + }), + name: 'grep____searchGitHub____mcp', + }, + id: 'grep____searchGitHub____mcp_0_6RnOMTF0', + thoughtSignature: 'test-signature', + type: 'function', + }, + ], + }, + { + content: '', + name: 'grep____searchGitHub____mcp', + role: 'tool', + tool_call_id: 'grep____searchGitHub____mcp_0_6RnOMTF0', + }, + ]; + + const contents = await buildGoogleMessages(messages); + + expect(contents).toEqual([ + { parts: [{ text: 'system prompt' }], role: 'user' }, + { parts: [{ text: 'LobeChat 最新版本' }], role: 'user' }, + { + parts: [ + { + functionCall: { + args: { + language: ['JSON'], + path: 'package.json', + query: '"version":', + repo: 'lobehub/lobe-chat', + }, + name: 'grep____searchGitHub____mcp', + }, + thoughtSignature: 'test-signature', + }, + ], + role: 'model', + }, + { + parts: [ + { + functionResponse: { + name: 'grep____searchGitHub____mcp', + response: { result: '' }, + }, + }, + ], + role: 'user', + }, + ]); + }); }); describe('buildGoogleTool', () => { diff --git a/packages/model-runtime/src/core/contextBuilders/google.ts b/packages/model-runtime/src/core/contextBuilders/google.ts index 9e885baa0f..18c61e4539 100644 --- a/packages/model-runtime/src/core/contextBuilders/google.ts +++ b/packages/model-runtime/src/core/contextBuilders/google.ts @@ -11,6 +11,12 @@ import { ChatCompletionTool, OpenAIChatMessage, UserMessageContentPart } from '. import { safeParseJSON } from '../../utils/safeParseJSON'; import { parseDataUri } from '../../utils/uriParser'; +/** + * Magic thoughtSignature + * @see https://ai.google.dev/gemini-api/docs/thought-signatures#model-behavior:~:text=context_engineering_is_the_way_to_go + */ +export const GEMINI_MAGIC_THOUGHT_SIGNATURE = 'context_engineering_is_the_way_to_go'; + /** * Convert OpenAI content part to Google Part format */ @@ -95,6 +101,7 @@ export const buildGoogleMessage = async ( args: safeParseJSON(tool.function.arguments)!, name: tool.function.name, }, + thoughtSignature: tool.thoughtSignature, })), role: 'model', }; @@ -155,7 +162,43 @@ export const buildGoogleMessages = async (messages: OpenAIChatMessage[]): Promis const contents = await Promise.all(pools); // Filter out empty messages: contents.parts must not be empty. - return contents.filter((content: Content) => content.parts && content.parts.length > 0); + const filteredContents = contents.filter( + (content: Content) => content.parts && content.parts.length > 0, + ); + + // Check if the last message is a tool message + const lastMessage = messages.at(-1); + const shouldAddMagicSignature = lastMessage?.role === 'tool'; + + if (shouldAddMagicSignature) { + // Find the last user message index in filtered contents + let lastUserIndex = -1; + for (let i = filteredContents.length - 1; i >= 0; i--) { + if (filteredContents[i].role === 'user') { + // Skip if it's a functionResponse (tool result) + const hasFunctionResponse = filteredContents[i].parts?.some((p) => p.functionResponse); + if (!hasFunctionResponse) { + lastUserIndex = i; + break; + } + } + } + + // Add magic signature to all function calls after last user message that don't have thoughtSignature + for (let i = lastUserIndex + 1; i < filteredContents.length; i++) { + const content = filteredContents[i]; + if (content.role === 'model' && content.parts) { + for (const part of content.parts) { + if (part.functionCall && !part.thoughtSignature) { + // Only add magic signature if thoughtSignature doesn't exist + part.thoughtSignature = GEMINI_MAGIC_THOUGHT_SIGNATURE; + } + } + } + } + } + + return filteredContents; }; /** diff --git a/packages/model-runtime/src/core/streams/google/google-ai.test.ts b/packages/model-runtime/src/core/streams/google/google-ai.test.ts index 41c14c7561..aea8dd54cc 100644 --- a/packages/model-runtime/src/core/streams/google/google-ai.test.ts +++ b/packages/model-runtime/src/core/streams/google/google-ai.test.ts @@ -4,938 +4,1239 @@ import { describe, expect, it, vi } from 'vitest'; import * as uuidModule from '../../../utils/uuid'; import { GoogleGenerativeAIStream, LOBE_ERROR_KEY } from './index'; +/** + * Helper function to decode stream chunks into string array + */ +async function decodeStreamChunks(stream: ReadableStream): Promise { + const decoder = new TextDecoder(); + const chunks: string[] = []; + + // @ts-ignore + for await (const chunk of stream) { + chunks.push(decoder.decode(chunk, { stream: true })); + } + + return chunks; +} + describe('GoogleGenerativeAIStream', () => { - it('should transform Google Generative AI stream to protocol stream', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1').mockReturnValueOnce('abcd1234'); + describe('Basic functionality', () => { + it('should transform Google Generative AI stream to protocol stream', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1').mockReturnValueOnce('abcd1234'); - const mockGenerateContentResponse = (text: string, functionCalls?: any[]) => - ({ - text: text, - functionCalls: functionCalls, - }) as unknown as GenerateContentResponse; + const mockGoogleStream = new ReadableStream({ + start(controller) { + // Text chunk + controller.enqueue({ + text: 'Hello', + candidates: [{ content: { parts: [{ text: 'Hello' }], role: 'model' }, index: 0 }], + } as unknown as GenerateContentResponse); - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue(mockGenerateContentResponse('Hello')); + // Function call chunk + controller.enqueue({ + text: '', + candidates: [ + { + content: { + parts: [ + { + functionCall: { + name: 'testFunction', + args: { arg1: 'value1' }, + }, + }, + ], + role: 'model', + }, + index: 0, + }, + ], + } as unknown as GenerateContentResponse); - controller.enqueue( - mockGenerateContentResponse('', [{ name: 'testFunction', args: { arg1: 'value1' } }]), - ); + // Final chunk with finishReason and usageMetadata to mark terminal event + controller.enqueue({ + text: ' world!', + candidates: [ + { + content: { parts: [{ text: ' world!' }], role: 'model' }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 1, + totalTokenCount: 1, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 1 }], + }, + modelVersion: 'gemini-test', + } as unknown as GenerateContentResponse); - // final chunk should include finishReason and usageMetadata to mark terminal event - controller.enqueue({ - text: ' world!', - candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], + controller.close(); + }, + }); + + const onStartMock = vi.fn(); + const onTextMock = vi.fn(); + const onToolCallMock = vi.fn(); + const onCompletionMock = vi.fn(); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream, { + callbacks: { + onStart: onStartMock, + onText: onTextMock, + onToolsCalling: onToolCallMock, + onCompletion: onCompletionMock, + }, + }); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + // text + 'id: chat_1\n', + 'event: text\n', + `data: "Hello"\n\n`, + + // tool call + 'id: chat_1\n', + 'event: tool_calls\n', + `data: [{"function":{"arguments":"{\\"arg1\\":\\"value1\\"}","name":"testFunction"},"id":"testFunction_0_abcd1234","index":0,"type":"function"}]\n\n`, + + // text + 'id: chat_1\n', + 'event: text\n', + `data: " world!"\n\n`, + // stop + 'id: chat_1\n', + 'event: stop\n', + `data: "STOP"\n\n`, + // usage + 'id: chat_1\n', + 'event: usage\n', + `data: {"inputTextTokens":1,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":1,"totalOutputTokens":0,"totalTokens":1}\n\n`, + ]); + + expect(onStartMock).toHaveBeenCalledTimes(1); + expect(onTextMock).toHaveBeenNthCalledWith(1, 'Hello'); + expect(onTextMock).toHaveBeenNthCalledWith(2, ' world!'); + expect(onToolCallMock).toHaveBeenCalledTimes(1); + expect(onCompletionMock).toHaveBeenCalledTimes(1); + }); + + it('should handle empty stream', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('E5M9dFKw'); + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue({ + candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], + usageMetadata: { + promptTokenCount: 0, + cachedContentTokenCount: 0, + totalTokenCount: 0, + promptTokensDetails: [ + { modality: 'TEXT', tokenCount: 0 }, + { modality: 'IMAGE', tokenCount: 0 }, + ], + }, + modelVersion: 'gemini-test', + } as unknown as GenerateContentResponse); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + 'id: chat_E5M9dFKw\n', + 'event: stop\n', + `data: "STOP"\n\n`, + 'id: chat_E5M9dFKw\n', + 'event: usage\n', + `data: {"inputCachedTokens":0,"inputImageTokens":0,"inputTextTokens":0,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":0,"totalOutputTokens":0,"totalTokens":0}\n\n`, + ]); + }); + + it('should return undefined data without text', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + candidates: [ + { + content: { parts: [{ text: '234' }], role: 'model' }, + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + text: '234', + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 3, + totalTokenCount: 122, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + thoughtsTokenCount: 100, + }, + modelVersion: 'gemini-2.5-flash-preview-04-17', + }, + { + text: '', + candidates: [ + { + content: { parts: [{ text: '' }], role: 'model' }, + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 3, + totalTokenCount: 122, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }], + thoughtsTokenCount: 100, + }, + modelVersion: 'gemini-2.5-flash-preview-04-17', + }, + { + text: '567890\n', + candidates: [ + { + content: { parts: [{ text: '567890\n' }], role: 'model' }, + finishReason: 'STOP', + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 11, + totalTokenCount: 131, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], + thoughtsTokenCount: 100, + }, + modelVersion: 'gemini-2.5-flash-preview-04-17', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: text', + 'data: "234"\n', + + 'id: chat_1', + 'event: text', + 'data: ""\n', + + 'id: chat_1', + 'event: text', + `data: "567890\\n"\n`, + // stop + 'id: chat_1', + 'event: stop', + `data: "STOP"\n`, + // usage + 'id: chat_1', + 'event: usage', + `data: {"inputTextTokens":19,"outputImageTokens":0,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`, + ].map((i) => i + '\n'), + ); + }); + }); + + describe('Reasoning and Thought', () => { + it('should handle thought candidate part', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + candidates: [ + { + content: { + parts: [{ text: '**Understanding the Conditional Logic**\n\n', thought: true }], + role: 'model', + }, + index: 0, + }, + ], + text: '**Understanding the Conditional Logic**\n\n', + usageMetadata: { + promptTokenCount: 38, + candidatesTokenCount: 7, + totalTokenCount: 301, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], + thoughtsTokenCount: 256, + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + { + candidates: [ + { + content: { + parts: [{ text: '**Finalizing Interpretation**\n\n', thought: true }], + role: 'model', + }, + index: 0, + }, + ], + text: '**Finalizing Interpretation**\n\n', + usageMetadata: { + promptTokenCount: 38, + candidatesTokenCount: 13, + totalTokenCount: 355, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], + thoughtsTokenCount: 304, + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + { + candidates: [ + { + content: { + parts: [{ text: '简单来说，' }], + role: 'model', + }, + index: 0, + }, + ], + text: '简单来说，', + usageMetadata: { + promptTokenCount: 38, + candidatesTokenCount: 16, + totalTokenCount: 358, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], + thoughtsTokenCount: 304, + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + { + candidates: [ + { + content: { parts: [{ text: '文本内容。' }], role: 'model' }, + finishReason: 'STOP', + index: 0, + }, + ], + text: '文本内容。', + usageMetadata: { + promptTokenCount: 38, + candidatesTokenCount: 19, + totalTokenCount: 361, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], + thoughtsTokenCount: 304, + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: reasoning', + 'data: "**Understanding the Conditional Logic**\\n\\n"\n', + + 'id: chat_1', + 'event: reasoning', + `data: "**Finalizing Interpretation**\\n\\n"\n`, + + 'id: chat_1', + 'event: text', + `data: "简单来说，"\n`, + + 'id: chat_1', + 'event: text', + `data: "文本内容。"\n`, + // stop + 'id: chat_1', + 'event: stop', + `data: "STOP"\n`, + // usage + 'id: chat_1', + 'event: usage', + `data: {"inputTextTokens":38,"outputImageTokens":0,"outputReasoningTokens":304,"outputTextTokens":19,"totalInputTokens":38,"totalOutputTokens":323,"totalTokens":361}\n`, + ].map((i) => i + '\n'), + ); + }); + + it('should handle stop with content and thought', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + candidates: [ + { + content: { parts: [{ text: '234' }], role: 'model' }, + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + text: '234', + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 3, + totalTokenCount: 122, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + thoughtsTokenCount: 100, + }, + modelVersion: 'gemini-2.5-flash-preview-04-17', + }, + { + text: '567890\n', + candidates: [ + { + content: { parts: [{ text: '567890\n' }], role: 'model' }, + finishReason: 'STOP', + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 11, + totalTokenCount: 131, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], + thoughtsTokenCount: 100, + }, + modelVersion: 'gemini-2.5-flash-preview-04-17', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: text', + 'data: "234"\n', + + 'id: chat_1', + 'event: text', + `data: "567890\\n"\n`, + // stop + 'id: chat_1', + 'event: stop', + `data: "STOP"\n`, + // usage + 'id: chat_1', + 'event: usage', + `data: {"inputTextTokens":19,"outputImageTokens":0,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`, + ].map((i) => i + '\n'), + ); + }); + }); + + describe('Usage and Token counting', () => { + it('should handle token count', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = { + candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], + usageMetadata: { + promptTokenCount: 266, + totalTokenCount: 266, + promptTokensDetails: [ + { modality: 'TEXT', tokenCount: 8 }, + { modality: 'IMAGE', tokenCount: 258 }, + ], + }, + modelVersion: 'gemini-2.0-flash-exp', + }; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue(data); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + // stop + 'id: chat_1\n', + 'event: stop\n', + `data: "STOP"\n\n`, + // usage + 'id: chat_1\n', + 'event: usage\n', + `data: {"inputImageTokens":258,"inputTextTokens":8,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":266,"totalOutputTokens":0,"totalTokens":266}\n\n`, + ]); + }); + + it('should handle token count with cached token count', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = { + candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], + usageMetadata: { + promptTokenCount: 15725, + candidatesTokenCount: 1053, + totalTokenCount: 16778, + cachedContentTokenCount: 14286, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 15725 }], + cacheTokensDetails: [{ modality: 'TEXT', tokenCount: 14286 }], + }, + modelVersion: 'gemini-2.0-flash-exp', + }; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue(data); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + // stop + 'id: chat_1\n', + 'event: stop\n', + `data: "STOP"\n\n`, + // usage + 'id: chat_1\n', + 'event: usage\n', + `data: {"inputCacheMissTokens":1439,"inputCachedTokens":14286,"inputTextTokens":15725,"outputImageTokens":0,"outputTextTokens":1053,"totalInputTokens":15725,"totalOutputTokens":1053,"totalTokens":16778}\n\n`, + ]); + }); + + it('should handle stop with content', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + candidates: [ + { + content: { parts: [{ text: '234' }], role: 'model' }, + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + text: '234', + usageMetadata: { + promptTokenCount: 20, + totalTokenCount: 20, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }], + }, + modelVersion: 'gemini-2.0-flash-exp-image-generation', + }, + { + text: '567890\n', + candidates: [ + { + content: { parts: [{ text: '567890\n' }], role: 'model' }, + finishReason: 'STOP', + safetyRatings: [ + { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, + { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, + ], + }, + ], + usageMetadata: { + promptTokenCount: 19, + candidatesTokenCount: 11, + totalTokenCount: 30, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], + candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], + }, + modelVersion: 'gemini-2.0-flash-exp-image-generation', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: text', + 'data: "234"\n', + + 'id: chat_1', + 'event: text', + `data: "567890\\n"\n`, + // stop + 'id: chat_1', + 'event: stop', + `data: "STOP"\n`, + // usage + 'id: chat_1', + 'event: usage', + `data: {"inputTextTokens":19,"outputImageTokens":0,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`, + ].map((i) => i + '\n'), + ); + }); + }); + + describe('Special content types', () => { + it('should handle image', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = { + candidates: [ + { + content: { + parts: [{ inlineData: { mimeType: 'image/png', data: 'iVBORw0KGgoAA' } }], + role: 'model', + }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 6, + totalTokenCount: 6, + promptTokensDetails: [ + { modality: 'TEXT', tokenCount: 6 }, + { modality: 'IMAGE', tokenCount: 0 }, + ], + }, + modelVersion: 'gemini-2.0-flash-exp', + }; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue(data); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + // image + 'id: chat_1\n', + 'event: base64_image\n', + `data: "data:image/png;base64,iVBORw0KGgoAA"\n\n`, + // stop + 'id: chat_1\n', + 'event: stop\n', + `data: "STOP"\n\n`, + // usage + 'id: chat_1\n', + 'event: usage\n', + `data: {"inputImageTokens":0,"inputTextTokens":6,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":6,"totalOutputTokens":0,"totalTokens":6}\n\n`, + ]); + }); + + it('should handle groundingMetadata', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + text: '123', + candidates: [ + { + content: { + parts: [ + { + text: '123', + }, + ], + role: 'model', + }, + index: 0, + groundingMetadata: {}, + }, + ], + usageMetadata: { + promptTokenCount: 9, + candidatesTokenCount: 18, + totalTokenCount: 27, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 9, + }, + ], + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + { + text: '45678', + candidates: [ + { + content: { + parts: [ + { + text: '45678', + }, + ], + role: 'model', + }, + finishReason: 'STOP', + index: 0, + groundingMetadata: { + searchEntryPoint: { + renderedContent: 'content\n', + }, + groundingChunks: [ + { + web: { + uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545', + title: 'npmjs.com', + }, + }, + { + web: { + uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334', + title: 'google.dev', + }, + }, + ], + groundingSupports: [ + { + segment: { + startIndex: 63, + endIndex: 67, + text: '1。', + }, + groundingChunkIndices: [0], + confidenceScores: [1], + }, + { + segment: { + startIndex: 69, + endIndex: 187, + text: 'SDK。', + }, + groundingChunkIndices: [1], + confidenceScores: [1], + }, + ], + webSearchQueries: ['sdk latest version'], + }, + }, + ], + usageMetadata: { + promptTokenCount: 9, + candidatesTokenCount: 122, + totalTokenCount: 131, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 9, + }, + ], + }, + modelVersion: 'models/gemini-2.5-flash-preview-04-17', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: text', + 'data: "123"\n', + + 'id: chat_1', + 'event: text', + 'data: "45678"\n', + + 'id: chat_1', + 'event: grounding', + `data: {\"citations\":[{\"favicon\":\"npmjs.com\",\"title\":\"npmjs.com\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545\"},{\"favicon\":\"google.dev\",\"title\":\"google.dev\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334\"}],\"searchQueries\":[\"sdk latest version\"]}\n`, + // stop + 'id: chat_1', + 'event: stop', + `data: "STOP"\n`, + // usage + 'id: chat_1', + 'event: usage', + `data: {"inputTextTokens":9,"outputImageTokens":0,"outputTextTokens":122,"totalInputTokens":9,"totalOutputTokens":122,"totalTokens":131}\n`, + ].map((i) => i + '\n'), + ); + }); + }); + + describe('Tool calls', () => { + it('should handle tool calls with thoughtSignature', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1').mockReturnValueOnce('abcd1234'); + + const data = [ + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + name: 'grep____searchGitHub____mcp', + args: { + query: '"version":', + repo: 'lobehub/lobe-chat', + path: 'package.json', + }, + }, + thoughtSignature: '123', + }, + ], + role: 'model', + }, + index: 0, + }, + ], + modelVersion: 'gemini-3-pro-preview', + responseId: 'UVcdaZ26ILac_uMP9ZOeiQ0', + usageMetadata: { + promptTokenCount: 1171, + candidatesTokenCount: 41, + totalTokenCount: 1408, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 1171 }], + thoughtsTokenCount: 196, + }, + }, + { + candidates: [ + { + content: { parts: [{ text: '' }], role: 'model' }, + finishReason: 'STOP', + index: 0, + }, + ], + modelVersion: 'gemini-3-pro-preview', + responseId: 'UVcdaZ26ILac_uMP9ZOeiQ0', + usageMetadata: { + promptTokenCount: 1171, + candidatesTokenCount: 41, + totalTokenCount: 1408, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 1171 }], + thoughtsTokenCount: 196, + }, + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: tool_calls', + 'data: [{"function":{"arguments":"{\\"query\\":\\"\\\\\\\"version\\\\\\":\\",\\"repo\\":\\"lobehub/lobe-chat\\",\\"path\\":\\"package.json\\"}","name":"grep____searchGitHub____mcp"},"id":"grep____searchGitHub____mcp_0_abcd1234","index":0,"thoughtSignature":"123","type":"function"}]\n', + + 'id: chat_1', + 'event: stop', + 'data: "STOP"\n', + + 'id: chat_1', + 'event: usage', + 'data: {"inputTextTokens":1171,"outputImageTokens":0,"outputReasoningTokens":196,"outputTextTokens":41,"totalInputTokens":1171,"totalOutputTokens":237,"totalTokens":1408}\n', + ].map((i) => i + '\n'), + ); + }); + + it('should handle parallel tool calls', async () => { + vi.spyOn(uuidModule, 'nanoid') + .mockReturnValueOnce('1') + .mockReturnValueOnce('abcd1234') + .mockReturnValueOnce('efgh5678'); + + const data = [ + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + name: 'get_current_temperature', + args: { + location: 'Paris', + }, + }, + thoughtSignature: 'ErEDCq4DAdHtim...', + }, + ], + role: 'model', + }, + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 72, + candidatesTokenCount: 18, + totalTokenCount: 167, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 72, + }, + ], + thoughtsTokenCount: 77, + }, + modelVersion: 'gemini-3-pro-preview', + responseId: 'UDcdaZviO4jojMcPycPDkQY', + }, + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + name: 'get_current_temperature', + args: { + location: 'London', + }, + }, + }, + ], + role: 'model', + }, + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 72, + candidatesTokenCount: 36, + totalTokenCount: 185, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 72, + }, + ], + thoughtsTokenCount: 77, + }, + modelVersion: 'gemini-3-pro-preview', + responseId: 'UDcdaZviO4jojMcPycPDkQY', + }, + { + candidates: [ + { + content: { + parts: [ + { + text: '', + }, + ], + role: 'model', + }, + finishReason: 'STOP', + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 72, + candidatesTokenCount: 36, + totalTokenCount: 185, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 72, + }, + ], + thoughtsTokenCount: 77, + }, + modelVersion: 'gemini-3-pro-preview', + responseId: 'UDcdaZviO4jojMcPycPDkQY', + }, + ]; + + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: tool_calls', + 'data: [{"function":{"arguments":"{\\"location\\":\\"Paris\\"}","name":"get_current_temperature"},"id":"get_current_temperature_0_abcd1234","index":0,"thoughtSignature":"ErEDCq4DAdHtim...","type":"function"}]\n', + + 'id: chat_1', + 'event: tool_calls', + 'data: [{"function":{"arguments":"{\\"location\\":\\"London\\"}","name":"get_current_temperature"},"id":"get_current_temperature_0_efgh5678","index":0,"type":"function"}]\n', + + 'id: chat_1', + 'event: stop', + 'data: "STOP"\n', + + 'id: chat_1', + 'event: usage', + 'data: {"inputTextTokens":72,"outputImageTokens":0,"outputReasoningTokens":77,"outputTextTokens":36,"totalInputTokens":72,"totalOutputTokens":113,"totalTokens":185}\n', + ].map((i) => i + '\n'), + ); + }); + + it('should handle thoughtSignature with empty text', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + + const data = [ + { + candidates: [ + { + content: { + parts: [ + { + text: '你好！很高兴为你服务。请问有什么我可以帮你的吗？\n\n无论是回答问题、协助写作、翻译，还是随便聊聊，我都随时待命！', + }, + ], + role: 'model', + }, + index: 0, + }, + ], usageMetadata: { promptTokenCount: 1, - totalTokenCount: 1, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 1 }], - }, - modelVersion: 'gemini-test', - } as unknown as GenerateContentResponse); - - controller.close(); - }, - }); - - const onStartMock = vi.fn(); - const onTextMock = vi.fn(); - const onToolCallMock = vi.fn(); - const onCompletionMock = vi.fn(); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream, { - callbacks: { - onStart: onStartMock, - onText: onTextMock, - onToolsCalling: onToolCallMock, - onCompletion: onCompletionMock, - }, - }); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - // text - 'id: chat_1\n', - 'event: text\n', - `data: "Hello"\n\n`, - - // tool call - 'id: chat_1\n', - 'event: tool_calls\n', - `data: [{"function":{"arguments":"{\\"arg1\\":\\"value1\\"}","name":"testFunction"},"id":"testFunction_0_abcd1234","index":0,"type":"function"}]\n\n`, - - // text - 'id: chat_1\n', - 'event: text\n', - `data: " world!"\n\n`, - // stop - 'id: chat_1\n', - 'event: stop\n', - `data: "STOP"\n\n`, - // usage - 'id: chat_1\n', - 'event: usage\n', - `data: {"inputTextTokens":1,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":1,"totalOutputTokens":0,"totalTokens":1}\n\n`, - ]); - - expect(onStartMock).toHaveBeenCalledTimes(1); - expect(onTextMock).toHaveBeenNthCalledWith(1, 'Hello'); - expect(onTextMock).toHaveBeenNthCalledWith(2, ' world!'); - expect(onToolCallMock).toHaveBeenCalledTimes(1); - expect(onCompletionMock).toHaveBeenCalledTimes(1); - }); - - it('should handle empty stream', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('E5M9dFKw'); - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue({ - candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], - usageMetadata: { - promptTokenCount: 0, - cachedContentTokenCount: 0, - totalTokenCount: 0, + candidatesTokenCount: 35, + totalTokenCount: 712, promptTokensDetails: [ - { modality: 'TEXT', tokenCount: 0 }, - { modality: 'IMAGE', tokenCount: 0 }, - ], - }, - modelVersion: 'gemini-test', - } as unknown as GenerateContentResponse); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - 'id: chat_E5M9dFKw\n', - 'event: stop\n', - `data: "STOP"\n\n`, - 'id: chat_E5M9dFKw\n', - 'event: usage\n', - `data: {"inputCachedTokens":0,"inputImageTokens":0,"inputTextTokens":0,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":0,"totalOutputTokens":0,"totalTokens":0}\n\n`, - ]); - }); - - it('should handle image', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = { - candidates: [ - { - content: { - parts: [{ inlineData: { mimeType: 'image/png', data: 'iVBORw0KGgoAA' } }], - role: 'model', - }, - finishReason: 'STOP', - index: 0, - }, - ], - usageMetadata: { - promptTokenCount: 6, - totalTokenCount: 6, - promptTokensDetails: [ - { modality: 'TEXT', tokenCount: 6 }, - { modality: 'IMAGE', tokenCount: 0 }, - ], - }, - modelVersion: 'gemini-2.0-flash-exp', - }; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue(data); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - // image - 'id: chat_1\n', - 'event: base64_image\n', - `data: "data:image/png;base64,iVBORw0KGgoAA"\n\n`, - // stop - 'id: chat_1\n', - 'event: stop\n', - `data: "STOP"\n\n`, - // usage - 'id: chat_1\n', - 'event: usage\n', - `data: {"inputImageTokens":0,"inputTextTokens":6,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":6,"totalOutputTokens":0,"totalTokens":6}\n\n`, - ]); - }); - - it('should handle token count', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = { - candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], - usageMetadata: { - promptTokenCount: 266, - totalTokenCount: 266, - promptTokensDetails: [ - { modality: 'TEXT', tokenCount: 8 }, - { modality: 'IMAGE', tokenCount: 258 }, - ], - }, - modelVersion: 'gemini-2.0-flash-exp', - }; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue(data); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - // stop - 'id: chat_1\n', - 'event: stop\n', - `data: "STOP"\n\n`, - // usage - 'id: chat_1\n', - 'event: usage\n', - `data: {"inputImageTokens":258,"inputTextTokens":8,"outputImageTokens":0,"outputTextTokens":0,"totalInputTokens":266,"totalOutputTokens":0,"totalTokens":266}\n\n`, - ]); - }); - - it('should handle token count with cached token count', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = { - candidates: [{ content: { role: 'model' }, finishReason: 'STOP', index: 0 }], - usageMetadata: { - promptTokenCount: 15725, - candidatesTokenCount: 1053, - totalTokenCount: 16778, - cachedContentTokenCount: 14286, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 15725 }], - cacheTokensDetails: [{ modality: 'TEXT', tokenCount: 14286 }], - }, - modelVersion: 'gemini-2.0-flash-exp', - }; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue(data); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - // stop - 'id: chat_1\n', - 'event: stop\n', - `data: "STOP"\n\n`, - // usage - 'id: chat_1\n', - 'event: usage\n', - `data: {"inputCacheMissTokens":1439,"inputCachedTokens":14286,"inputTextTokens":15725,"outputImageTokens":0,"outputTextTokens":1053,"totalInputTokens":15725,"totalOutputTokens":1053,"totalTokens":16778}\n\n`, - ]); - }); - - it('should handle stop with content', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = [ - { - candidates: [ - { - content: { parts: [{ text: '234' }], role: 'model' }, - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - text: '234', - usageMetadata: { - promptTokenCount: 20, - totalTokenCount: 20, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }], - }, - modelVersion: 'gemini-2.0-flash-exp-image-generation', - }, - { - text: '567890\n', - candidates: [ - { - content: { parts: [{ text: '567890\n' }], role: 'model' }, - finishReason: 'STOP', - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 11, - totalTokenCount: 30, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], - }, - modelVersion: 'gemini-2.0-flash-exp-image-generation', - }, - ]; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - data.forEach((item) => { - controller.enqueue(item); - }); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual( - [ - 'id: chat_1', - 'event: text', - 'data: "234"\n', - - 'id: chat_1', - 'event: text', - `data: "567890\\n"\n`, - // stop - 'id: chat_1', - 'event: stop', - `data: "STOP"\n`, - // usage - 'id: chat_1', - 'event: usage', - `data: {"inputTextTokens":19,"outputImageTokens":0,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`, - ].map((i) => i + '\n'), - ); - }); - - it('should handle stop with content and thought', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = [ - { - candidates: [ - { - content: { parts: [{ text: '234' }], role: 'model' }, - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - text: '234', - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 3, - totalTokenCount: 122, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - thoughtsTokenCount: 100, - }, - modelVersion: 'gemini-2.5-flash-preview-04-17', - }, - { - text: '567890\n', - candidates: [ - { - content: { parts: [{ text: '567890\n' }], role: 'model' }, - finishReason: 'STOP', - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 11, - totalTokenCount: 131, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], - thoughtsTokenCount: 100, - }, - modelVersion: 'gemini-2.5-flash-preview-04-17', - }, - ]; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - data.forEach((item) => { - controller.enqueue(item); - }); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual( - [ - 'id: chat_1', - 'event: text', - 'data: "234"\n', - - 'id: chat_1', - 'event: text', - `data: "567890\\n"\n`, - // stop - 'id: chat_1', - 'event: stop', - `data: "STOP"\n`, - // usage - 'id: chat_1', - 'event: usage', - `data: {"inputTextTokens":19,"outputImageTokens":0,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`, - ].map((i) => i + '\n'), - ); - }); - - it('should handle thought candidate part', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = [ - { - candidates: [ - { - content: { - parts: [{ text: '**Understanding the Conditional Logic**\n\n', thought: true }], - role: 'model', - }, - index: 0, - }, - ], - text: '**Understanding the Conditional Logic**\n\n', - usageMetadata: { - promptTokenCount: 38, - candidatesTokenCount: 7, - totalTokenCount: 301, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], - thoughtsTokenCount: 256, - }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - { - candidates: [ - { - content: { - parts: [{ text: '**Finalizing Interpretation**\n\n', thought: true }], - role: 'model', - }, - index: 0, - }, - ], - text: '**Finalizing Interpretation**\n\n', - usageMetadata: { - promptTokenCount: 38, - candidatesTokenCount: 13, - totalTokenCount: 355, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], - thoughtsTokenCount: 304, - }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - { - candidates: [ - { - content: { - parts: [{ text: '简单来说，' }], - role: 'model', - }, - index: 0, - }, - ], - text: '简单来说，', - usageMetadata: { - promptTokenCount: 38, - candidatesTokenCount: 16, - totalTokenCount: 358, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], - thoughtsTokenCount: 304, - }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - { - candidates: [ - { - content: { parts: [{ text: '文本内容。' }], role: 'model' }, - finishReason: 'STOP', - index: 0, - }, - ], - text: '文本内容。', - usageMetadata: { - promptTokenCount: 38, - candidatesTokenCount: 19, - totalTokenCount: 361, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 38 }], - thoughtsTokenCount: 304, - }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - ]; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - data.forEach((item) => { - controller.enqueue(item); - }); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual( - [ - 'id: chat_1', - 'event: reasoning', - 'data: "**Understanding the Conditional Logic**\\n\\n"\n', - - 'id: chat_1', - 'event: reasoning', - `data: "**Finalizing Interpretation**\\n\\n"\n`, - - 'id: chat_1', - 'event: text', - `data: "简单来说，"\n`, - - 'id: chat_1', - 'event: text', - `data: "文本内容。"\n`, - // stop - 'id: chat_1', - 'event: stop', - `data: "STOP"\n`, - // usage - 'id: chat_1', - 'event: usage', - `data: {"inputTextTokens":38,"outputImageTokens":0,"outputReasoningTokens":304,"outputTextTokens":19,"totalInputTokens":38,"totalOutputTokens":323,"totalTokens":361}\n`, - ].map((i) => i + '\n'), - ); - }); - - it('should return undefined data without text', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = [ - { - candidates: [ - { - content: { parts: [{ text: '234' }], role: 'model' }, - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - text: '234', - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 3, - totalTokenCount: 122, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - thoughtsTokenCount: 100, - }, - modelVersion: 'gemini-2.5-flash-preview-04-17', - }, - { - text: '', - candidates: [ - { - content: { parts: [{ text: '' }], role: 'model' }, - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 3, - totalTokenCount: 122, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 3 }], - thoughtsTokenCount: 100, - }, - modelVersion: 'gemini-2.5-flash-preview-04-17', - }, - { - text: '567890\n', - candidates: [ - { - content: { parts: [{ text: '567890\n' }], role: 'model' }, - finishReason: 'STOP', - safetyRatings: [ - { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' }, - { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' }, - ], - }, - ], - usageMetadata: { - promptTokenCount: 19, - candidatesTokenCount: 11, - totalTokenCount: 131, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }], - candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }], - thoughtsTokenCount: 100, - }, - modelVersion: 'gemini-2.5-flash-preview-04-17', - }, - ]; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - data.forEach((item) => { - controller.enqueue(item); - }); - - controller.close(); - }, - }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual( - [ - 'id: chat_1', - 'event: text', - 'data: "234"\n', - - 'id: chat_1', - 'event: text', - 'data: ""\n', - - 'id: chat_1', - 'event: text', - `data: "567890\\n"\n`, - // stop - 'id: chat_1', - 'event: stop', - `data: "STOP"\n`, - // usage - 'id: chat_1', - 'event: usage', - `data: {"inputTextTokens":19,"outputImageTokens":0,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`, - ].map((i) => i + '\n'), - ); - }); - - it('should handle groundingMetadata', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - - const data = [ - { - text: '123', - candidates: [ - { - content: { - parts: [ - { - text: '123', - }, - ], - role: 'model', - }, - index: 0, - groundingMetadata: {}, - }, - ], - usageMetadata: { - promptTokenCount: 9, - candidatesTokenCount: 18, - totalTokenCount: 27, - promptTokensDetails: [ - { - modality: 'TEXT', - tokenCount: 9, - }, - ], - }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - { - text: '45678', - candidates: [ - { - content: { - parts: [ - { - text: '45678', - }, - ], - role: 'model', - }, - finishReason: 'STOP', - index: 0, - groundingMetadata: { - searchEntryPoint: { - renderedContent: 'content\n', + { + modality: 'TEXT', + tokenCount: 1, }, - groundingChunks: [ - { - web: { - uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545', - title: 'npmjs.com', - }, - }, - { - web: { - uri: 'https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334', - title: 'google.dev', - }, - }, - ], - groundingSupports: [ - { - segment: { - startIndex: 63, - endIndex: 67, - text: '1。', - }, - groundingChunkIndices: [0], - confidenceScores: [1], - }, - { - segment: { - startIndex: 69, - endIndex: 187, - text: 'SDK。', - }, - groundingChunkIndices: [1], - confidenceScores: [1], - }, - ], - webSearchQueries: ['sdk latest version'], - }, + ], + thoughtsTokenCount: 676, }, - ], - usageMetadata: { - promptTokenCount: 9, - candidatesTokenCount: 122, - totalTokenCount: 131, - promptTokensDetails: [ + modelVersion: 'gemini-3-pro-preview', + responseId: 'lTcdaf_1FrONjMcP24Sz6QQ', + }, + { + candidates: [ { - modality: 'TEXT', - tokenCount: 9, + content: { + parts: [ + { + text: '', + thoughtSignature: 'Ep8YCpwYAdHtim...', + }, + ], + role: 'model', + }, + finishReason: 'STOP', + index: 0, }, ], + usageMetadata: { + promptTokenCount: 1, + candidatesTokenCount: 35, + totalTokenCount: 712, + promptTokensDetails: [ + { + modality: 'TEXT', + tokenCount: 1, + }, + ], + thoughtsTokenCount: 676, + }, + modelVersion: 'gemini-3-pro-preview', + responseId: 'lTcdaf_1FrONjMcP24Sz6QQ', }, - modelVersion: 'models/gemini-2.5-flash-preview-04-17', - }, - ]; + ]; - const mockGoogleStream = new ReadableStream({ - start(controller) { - data.forEach((item) => { - controller.enqueue(item); - }); + const mockGoogleStream = new ReadableStream({ + start(controller) { + data.forEach((item) => { + controller.enqueue(item); + }); - controller.close(); - }, + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual( + [ + 'id: chat_1', + 'event: text', + 'data: "你好！很高兴为你服务。请问有什么我可以帮你的吗？\\n\\n无论是回答问题、协助写作、翻译，还是随便聊聊，我都随时待命！"\n', + + 'id: chat_1', + 'event: stop', + 'data: "STOP"\n', + + 'id: chat_1', + 'event: usage', + 'data: {"inputTextTokens":1,"outputImageTokens":0,"outputReasoningTokens":676,"outputTextTokens":35,"totalInputTokens":1,"totalOutputTokens":711,"totalTokens":712}\n', + ].map((i) => i + '\n'), + ); }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual( - [ - 'id: chat_1', - 'event: text', - 'data: "123"\n', - - 'id: chat_1', - 'event: text', - 'data: "45678"\n', - - 'id: chat_1', - 'event: grounding', - `data: {\"citations\":[{\"favicon\":\"npmjs.com\",\"title\":\"npmjs.com\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXG1234545\"},{\"favicon\":\"google.dev\",\"title\":\"google.dev\",\"url\":\"https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXE9288334\"}],\"searchQueries\":[\"sdk latest version\"]}\n`, - // stop - 'id: chat_1', - 'event: stop', - `data: "STOP"\n`, - // usage - 'id: chat_1', - 'event: usage', - `data: {"inputTextTokens":9,"outputImageTokens":0,"outputTextTokens":122,"totalInputTokens":9,"totalOutputTokens":122,"totalTokens":131}\n`, - ].map((i) => i + '\n'), - ); }); - it('should handle promptFeedback with blockReason (PROHIBITED_CONTENT)', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + describe('Error handling', () => { + it('should handle promptFeedback with blockReason (PROHIBITED_CONTENT)', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - const data = { - promptFeedback: { - blockReason: 'PROHIBITED_CONTENT', - }, - usageMetadata: { - promptTokenCount: 4438, - totalTokenCount: 4438, - promptTokensDetails: [{ modality: 'TEXT', tokenCount: 4438 }], - }, - modelVersion: 'gemini-2.5-pro', - responseId: 'THOUaKaNOeiGz7IPjL_VgQc', - }; + const data = { + promptFeedback: { + blockReason: 'PROHIBITED_CONTENT', + }, + usageMetadata: { + promptTokenCount: 4438, + totalTokenCount: 4438, + promptTokensDetails: [{ modality: 'TEXT', tokenCount: 4438 }], + }, + modelVersion: 'gemini-2.5-pro', + responseId: 'THOUaKaNOeiGz7IPjL_VgQc', + }; - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue(data); - controller.close(); - }, + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue(data); + controller.close(); + }, + }); + + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + + const chunks = await decodeStreamChunks(protocolStream); + + expect(chunks).toEqual([ + 'id: chat_1\n', + 'event: error\n', + `data: {"body":{"context":{"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"}},"message":"Your request may contain prohibited content. Please adjust your request to comply with the usage guidelines.","provider":"google"},"type":"ProviderBizError"}\n\n`, + ]); }); - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); + it('should pass through injected lobe error marker', async () => { + vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); - const decoder = new TextDecoder(); - const chunks = []; + const errorPayload = { message: 'internal error', code: 123 }; - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } + const mockGoogleStream = new ReadableStream({ + start(controller) { + controller.enqueue({ [LOBE_ERROR_KEY]: errorPayload }); + controller.close(); + }, + }); - expect(chunks).toEqual([ - 'id: chat_1\n', - 'event: error\n', - `data: {"body":{"context":{"promptFeedback":{"blockReason":"PROHIBITED_CONTENT"}},"message":"Your request may contain prohibited content. Please adjust your request to comply with the usage guidelines.","provider":"google"},"type":"ProviderBizError"}\n\n`, - ]); - }); + const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - it('should pass through injected lobe error marker', async () => { - vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1'); + const chunks = await decodeStreamChunks(protocolStream); - const errorPayload = { message: 'internal error', code: 123 }; - - const mockGoogleStream = new ReadableStream({ - start(controller) { - controller.enqueue({ [LOBE_ERROR_KEY]: errorPayload }); - controller.close(); - }, + expect(chunks).toEqual([ + 'id: chat_1\n', + 'event: error\n', + `data: ${JSON.stringify(errorPayload)}\n\n`, + ]); }); - - const protocolStream = GoogleGenerativeAIStream(mockGoogleStream); - - const decoder = new TextDecoder(); - const chunks = []; - - // @ts-ignore - for await (const chunk of protocolStream) { - chunks.push(decoder.decode(chunk, { stream: true })); - } - - expect(chunks).toEqual([ - 'id: chat_1\n', - 'event: error\n', - `data: ${JSON.stringify(errorPayload)}\n\n`, - ]); }); }); diff --git a/packages/model-runtime/src/core/streams/google/index.ts b/packages/model-runtime/src/core/streams/google/index.ts index 0243f3c657..b02b6bd278 100644 --- a/packages/model-runtime/src/core/streams/google/index.ts +++ b/packages/model-runtime/src/core/streams/google/index.ts @@ -1,4 +1,4 @@ -import { GenerateContentResponse } from '@google/genai'; +import { GenerateContentResponse, Part } from '@google/genai'; import { GroundingSearch } from '@lobechat/types'; import { ChatStreamCallbacks } from '../../../types'; @@ -74,19 +74,27 @@ const transformGoogleGenerativeAIStream = ( } } - const functionCalls = chunk.functionCalls; + // Parse function calls from candidate.content.parts + const functionCalls = + candidate?.content?.parts + ?.filter((part: any) => part.functionCall) + .map((part: Part) => ({ + ...part.functionCall, + thoughtSignature: part.thoughtSignature, + })) || []; - if (functionCalls) { + if (functionCalls.length > 0) { return [ { data: functionCalls.map( - (value, index): StreamToolCallChunkData => ({ + (value, index: number): StreamToolCallChunkData => ({ function: { arguments: JSON.stringify(value.args), name: value.name, }, id: generateToolCallId(index, value.name), index: index, + thoughtSignature: value.thoughtSignature, type: 'function', }), ), @@ -97,7 +105,13 @@ const transformGoogleGenerativeAIStream = ( ]; } - const text = chunk.text; + // Parse text from candidate.content.parts + // Filter out thought content (thought: true) and thoughtSignature + const text = + candidate?.content?.parts + ?.filter((part: any) => part.text && !part.thought && !part.thoughtSignature) + .map((part: any) => part.text) + .join('') || ''; if (candidate) { // 首先检查是否为 reasoning 内容 (thought: true) diff --git a/packages/model-runtime/src/core/streams/protocol.ts b/packages/model-runtime/src/core/streams/protocol.ts index aa6705d181..01fe864d5a 100644 --- a/packages/model-runtime/src/core/streams/protocol.ts +++ b/packages/model-runtime/src/core/streams/protocol.ts @@ -98,6 +98,7 @@ export interface StreamToolCallChunkData { }; id?: string; index: number; + thoughtSignature?: string; type: 'function' | string; } diff --git a/packages/model-runtime/src/providers/google/index.test.ts b/packages/model-runtime/src/providers/google/index.test.ts index d0a88e213b..8566dac3e3 100644 --- a/packages/model-runtime/src/providers/google/index.test.ts +++ b/packages/model-runtime/src/providers/google/index.test.ts @@ -1,5 +1,5 @@ // @vitest-environment node -import { GenerateContentResponse, Tool } from '@google/genai'; +import { GenerateContentResponse } from '@google/genai'; import OpenAI from 'openai'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; diff --git a/packages/model-runtime/src/providers/google/index.ts b/packages/model-runtime/src/providers/google/index.ts index e4ffae69e1..c6db404943 100644 --- a/packages/model-runtime/src/providers/google/index.ts +++ b/packages/model-runtime/src/providers/google/index.ts @@ -267,19 +267,21 @@ export class LobeGoogleAI implements LobeRuntimeAI { const inputStartAt = Date.now(); - const geminiStreamResponse = await this.client.models.generateContentStream({ - config, - contents, - model, - }); - - const googleStream = this.createEnhancedStream(geminiStreamResponse, controller.signal); - const [prod, useForDebug] = googleStream.tee(); - + const finalPayload = { config, contents, model }; const key = this.isVertexAi ? 'DEBUG_VERTEX_AI_CHAT_COMPLETION' : 'DEBUG_GOOGLE_CHAT_COMPLETION'; + if (process.env[key] === '1') { + console.log('[requestPayload]'); + console.log(JSON.stringify(finalPayload), '\n'); + } + + const geminiStreamResponse = await this.client.models.generateContentStream(finalPayload); + + const googleStream = this.createEnhancedStream(geminiStreamResponse, controller.signal); + const [prod, useForDebug] = googleStream.tee(); + if (process.env[key] === '1') { debugStream(useForDebug).catch(); } diff --git a/packages/model-runtime/src/types/toolsCalling.ts b/packages/model-runtime/src/types/toolsCalling.ts index d5a711908a..9198f7b62b 100644 --- a/packages/model-runtime/src/types/toolsCalling.ts +++ b/packages/model-runtime/src/types/toolsCalling.ts @@ -1,5 +1,5 @@ -import { z } from 'zod'; import type { PartialDeep } from 'type-fest'; +import { z } from 'zod'; /** * The function that the model called. @@ -30,6 +30,7 @@ export interface MessageToolCall { */ id: string; + thoughtSignature?: string; /** * The type of the tool. Currently, only `function` is supported. */ @@ -42,6 +43,7 @@ export const MessageToolCallSchema = z.object({ name: z.string(), }), id: z.string(), + thoughtSignature: z.string().optional(), type: z.string(), }); diff --git a/packages/types/src/message/common/tools.ts b/packages/types/src/message/common/tools.ts index e55fcc2a2f..6b36bb4648 100644 --- a/packages/types/src/message/common/tools.ts +++ b/packages/types/src/message/common/tools.ts @@ -30,6 +30,7 @@ export interface ChatToolPayload { identifier: string; intervention?: ToolIntervention; result_msg_id?: string; + thoughtSignature?: string; type: LobeToolRenderType; } @@ -84,6 +85,7 @@ export interface MessageToolCall { */ id: string; + thoughtSignature?: string; /** * The type of the tool. Currently, only `function` is supported. */ @@ -108,6 +110,7 @@ export const ChatToolPayloadSchema = z.object({ identifier: z.string(), intervention: ToolInterventionSchema.optional(), result_msg_id: z.string().optional(), + thoughtSignature: z.string().optional(), type: z.string(), }); diff --git a/src/features/Conversation/Messages/Group/Error/index.tsx b/src/features/Conversation/Messages/Group/Error/index.tsx index a89542a849..57585908d1 100644 --- a/src/features/Conversation/Messages/Group/Error/index.tsx +++ b/src/features/Conversation/Messages/Group/Error/index.tsx @@ -15,12 +15,13 @@ export interface ErrorContentProps { const ErrorContent = memo(({ error, id }) => { const { t } = useTranslation('common'); - const errorProps = useErrorContent(error); const [deleteMessage] = useChatStore((s) => [s.deleteDBMessage]); const message = ; - if (!error?.message) { + const errorProps = useErrorContent(error); + + if (!errorProps?.message) { if (!message) return null; return {message}; } diff --git a/src/features/Conversation/Messages/Group/GroupItem.tsx b/src/features/Conversation/Messages/Group/GroupItem.tsx index 35458670be..a9f80220aa 100644 --- a/src/features/Conversation/Messages/Group/GroupItem.tsx +++ b/src/features/Conversation/Messages/Group/GroupItem.tsx @@ -30,10 +30,10 @@ const GroupItem = memo( }); }} > - + ) : ( - + ); }, isEqual, diff --git a/src/store/chat/slices/aiChat/actions/streamingExecutor.ts b/src/store/chat/slices/aiChat/actions/streamingExecutor.ts index 3160b7faab..7b61598215 100644 --- a/src/store/chat/slices/aiChat/actions/streamingExecutor.ts +++ b/src/store/chat/slices/aiChat/actions/streamingExecutor.ts @@ -76,7 +76,7 @@ export interface StreamingExecutorAction { tool_calls?: MessageToolCall[]; content: string; traceId?: string; - finishType?: 'done' | 'error' | 'abort'; + finishType?: string; usage?: ModelUsage; }>; /** @@ -283,13 +283,13 @@ export const streamingExecutor: StateCreator< let thinkingStartAt: number; let duration: number | undefined; let reasoningOperationId: string | undefined; - let finishType: 'done' | 'error' | 'abort' | undefined; + let finishType: string | undefined; // to upload image const uploadTasks: Map> = new Map(); // Throttle tool_calls updates to prevent excessive re-renders (max once per 300ms) const throttledUpdateToolCalls = throttle( - (toolCalls: any[]) => { + (toolCalls: MessageToolCall[]) => { internal_dispatchMessage( { id: messageId, @@ -366,7 +366,6 @@ export const streamingExecutor: StateCreator< throttledUpdateToolCalls.flush(); internal_toggleToolCallingStreaming(messageId, undefined); - tools = get().internal_transformToolCalls(parsedToolCalls); tool_calls = toolCalls; parsedToolCalls = parsedToolCalls.map((item) => ({ @@ -377,6 +376,8 @@ export const streamingExecutor: StateCreator< }, })); + tools = get().internal_transformToolCalls(parsedToolCalls); + isFunctionCall = true; } @@ -395,7 +396,7 @@ export const streamingExecutor: StateCreator< messageId, content, { - toolCalls: parsedToolCalls, + tools, reasoning: !!reasoning ? { ...reasoning, duration: duration && !isNaN(duration) ? duration : undefined } : undefined, diff --git a/src/store/chat/slices/message/actions/optimisticUpdate.ts b/src/store/chat/slices/message/actions/optimisticUpdate.ts index b8d44f8386..b953026460 100644 --- a/src/store/chat/slices/message/actions/optimisticUpdate.ts +++ b/src/store/chat/slices/message/actions/optimisticUpdate.ts @@ -3,11 +3,11 @@ import { ChatImageItem, ChatMessageError, ChatMessagePluginError, + ChatToolPayload, CreateMessageParams, GroundingSearch, MessageMetadata, MessagePluginItem, - MessageToolCall, ModelReasoning, UIChatMessage, UpdateMessageRAGParams, @@ -69,7 +69,7 @@ export interface MessageOptimisticUpdateAction { provider?: string; reasoning?: ModelReasoning; search?: GroundingSearch; - toolCalls?: MessageToolCall[]; + tools?: ChatToolPayload[]; }, context?: OptimisticUpdateContext, ) => Promise; @@ -204,22 +204,17 @@ export const messageOptimisticUpdate: StateCreator< }, optimisticUpdateMessageContent: async (id, content, extra, context) => { - const { - internal_dispatchMessage, - refreshMessages, - internal_transformToolCalls, - replaceMessages, - } = get(); + const { internal_dispatchMessage, refreshMessages, replaceMessages } = get(); // Due to the async update method and refresh need about 100ms // we need to update the message content at the frontend to avoid the update flick // refs: https://medium.com/@kyledeguzmanx/what-are-optimistic-updates-483662c3e171 - if (extra?.toolCalls) { + if (extra?.tools) { internal_dispatchMessage( { id, type: 'updateMessage', - value: { tools: internal_transformToolCalls(extra?.toolCalls) }, + value: { tools: extra?.tools }, }, context, ); @@ -246,7 +241,7 @@ export const messageOptimisticUpdate: StateCreator< provider: extra?.provider, reasoning: extra?.reasoning, search: extra?.search, - tools: extra?.toolCalls ? internal_transformToolCalls(extra?.toolCalls) : undefined, + tools: extra?.tools, }, { sessionId, topicId }, ); diff --git a/src/store/chat/slices/plugin/actions/internals.ts b/src/store/chat/slices/plugin/actions/internals.ts index 76e28463d0..61bb590b52 100644 --- a/src/store/chat/slices/plugin/actions/internals.ts +++ b/src/store/chat/slices/plugin/actions/internals.ts @@ -1,6 +1,6 @@ /* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */ import { ToolNameResolver } from '@lobechat/context-engine'; -import { MessageToolCall, ToolsCallingContext } from '@lobechat/types'; +import { ChatToolPayload, MessageToolCall, ToolsCallingContext } from '@lobechat/types'; import { LobeChatPluginManifest } from '@lobehub/chat-plugin-sdk'; import { StateCreator } from 'zustand/vanilla'; @@ -19,7 +19,7 @@ export interface PluginInternalsAction { /** * Transform tool calls from runtime format to storage format */ - internal_transformToolCalls: (toolCalls: MessageToolCall[]) => any[]; + internal_transformToolCalls: (toolCalls: MessageToolCall[]) => ChatToolPayload[]; /** * Construct tools calling context for plugin invocation