diff --git a/packages/context-engine/src/index.ts b/packages/context-engine/src/index.ts index a7414e1161b..366d637b4f2 100644 --- a/packages/context-engine/src/index.ts +++ b/packages/context-engine/src/index.ts @@ -19,6 +19,7 @@ export { // Processors export { + GroupMessageFlattenProcessor, HistoryTruncateProcessor, InputTemplateProcessor, MessageCleanupProcessor, diff --git a/packages/context-engine/src/processors/GroupMessageFlatten.ts b/packages/context-engine/src/processors/GroupMessageFlatten.ts new file mode 100644 index 00000000000..ae39d6e6e46 --- /dev/null +++ b/packages/context-engine/src/processors/GroupMessageFlatten.ts @@ -0,0 +1,144 @@ +import debug from 'debug'; + +import { BaseProcessor } from '../base/BaseProcessor'; +import type { PipelineContext, ProcessorOptions } from '../types'; + +const log = debug('context-engine:processor:GroupMessageFlattenProcessor'); + +/** + * Group Message Flatten Processor + * Responsible for flattening role=group messages into standard assistant + tool message sequences + * + * Group messages are created when assistant messages with tools are merged with their tool results. + * This processor converts them back to a flat structure that AI models can understand. + */ +export class GroupMessageFlattenProcessor extends BaseProcessor { + readonly name = 'GroupMessageFlattenProcessor'; + + constructor(options: ProcessorOptions = {}) { + super(options); + } + + protected async doProcess(context: PipelineContext): Promise { + const clonedContext = this.cloneContext(context); + + let processedCount = 0; + let groupMessagesFlattened = 0; + let assistantMessagesCreated = 0; + let toolMessagesCreated = 0; + + const newMessages: any[] = []; + + // Process each message + for (const message of clonedContext.messages) { + // Check if this is a group message with children field + if (message.role === 'group' && message.children) { + // If children array is empty, skip this message entirely (no content to flatten) + if (message.children.length === 0) { + continue; + } + + processedCount++; + groupMessagesFlattened++; + + log(`Flattening group message ${message.id} with ${message.children.length} children`); + + // Flatten each child + for (const child of message.children) { + // 1. Create assistant message from child + const assistantMsg: any = { + content: child.content || '', + createdAt: message.createdAt, + id: child.id, + meta: message.meta, + role: 'assistant', + updatedAt: message.updatedAt, + }; + + // Add tools if present (excluding result, which will be separate tool messages) + if (child.tools && child.tools.length > 0) { + assistantMsg.tools = child.tools.map((tool: any) => ({ + apiName: tool.apiName, + arguments: tool.arguments, + id: tool.id, + identifier: tool.identifier, + type: tool.type, + })); + } + + // Add reasoning if present (for models that support reasoning) + if (message.reasoning) { + assistantMsg.reasoning = message.reasoning; + } + + // Preserve other fields that might be needed + if (message.parentId) assistantMsg.parentId = message.parentId; + if (message.threadId) assistantMsg.threadId = message.threadId; + if (message.groupId) assistantMsg.groupId = message.groupId; + if (message.agentId) assistantMsg.agentId = message.agentId; + if (message.targetId) assistantMsg.targetId = message.targetId; + if (message.topicId) assistantMsg.topicId = message.topicId; + + newMessages.push(assistantMsg); + assistantMessagesCreated++; + + log(`Created assistant message ${assistantMsg.id} from child`); + + // 2. Create tool messages for each tool that has a result + if (child.tools) { + for (const tool of child.tools) { + if (tool.result) { + const toolMsg: any = { + content: tool.result.content, + createdAt: message.createdAt, + id: tool.result.id, + meta: message.meta, + plugin: { + apiName: tool.apiName, + arguments: tool.arguments, + id: tool.id, + identifier: tool.identifier, + type: tool.type, + }, + pluginError: tool.result.error || undefined, + pluginState: tool.result.state || undefined, + role: 'tool', + tool_call_id: tool.id, + updatedAt: message.updatedAt, + }; + + // Preserve parent message references + if (message.parentId) toolMsg.parentId = message.parentId; + if (message.threadId) toolMsg.threadId = message.threadId; + if (message.groupId) toolMsg.groupId = message.groupId; + if (message.topicId) toolMsg.topicId = message.topicId; + + newMessages.push(toolMsg); + toolMessagesCreated++; + + log(`Created tool message ${toolMsg.id} for tool call ${tool.id}`); + } + } + } + } + } else { + // Non-group message, keep as-is + newMessages.push(message); + } + } + + clonedContext.messages = newMessages; + + // Update metadata + clonedContext.metadata.groupMessagesFlattenProcessed = processedCount; + clonedContext.metadata.groupMessagesFlattened = groupMessagesFlattened; + clonedContext.metadata.assistantMessagesCreated = assistantMessagesCreated; + clonedContext.metadata.toolMessagesCreated = toolMessagesCreated; + + log( + `Group message flatten processing completed: ${groupMessagesFlattened} groups flattened, ${assistantMessagesCreated} assistant messages created, ${toolMessagesCreated} tool messages created`, + ); + + return this.markAsExecuted(clonedContext); + } +} diff --git a/packages/context-engine/src/processors/__tests__/GroupMessageFlatten.test.ts b/packages/context-engine/src/processors/__tests__/GroupMessageFlatten.test.ts new file mode 100644 index 00000000000..eb3ad243cf5 --- /dev/null +++ b/packages/context-engine/src/processors/__tests__/GroupMessageFlatten.test.ts @@ -0,0 +1,522 @@ +import { describe, expect, it } from 'vitest'; + +import { PipelineContext } from '../../types'; +import { GroupMessageFlattenProcessor } from '../GroupMessageFlatten'; + +describe('GroupMessageFlattenProcessor', () => { + const createContext = (messages: any[]): PipelineContext => ({ + initialState: { messages: [] }, + isAborted: false, + messages, + metadata: {}, + }); + + describe('Basic Scenarios', () => { + it('should flatten group message with single child and single tool result', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + createdAt: '2025-10-27T10:00:00.000Z', + updatedAt: '2025-10-27T10:00:10.000Z', + meta: { title: 'Test Agent' }, + children: [ + { + id: 'msg-1', + content: 'Checking weather', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{"query":"weather"}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-1', + content: 'Weather is sunny, 25°C', + error: null, + state: { cached: true }, + }, + }, + ], + usage: { totalTokens: 100 }, + performance: { tps: 20 }, + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Should create 2 messages: 1 assistant + 1 tool + expect(result.messages).toHaveLength(2); + + // Check assistant message + const assistantMsg = result.messages[0]; + expect(assistantMsg.role).toBe('assistant'); + expect(assistantMsg.id).toBe('msg-1'); + expect(assistantMsg.content).toBe('Checking weather'); + expect(assistantMsg.tools).toHaveLength(1); + expect(assistantMsg.tools[0]).toEqual({ + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{"query":"weather"}', + identifier: 'web-browsing', + }); + expect(assistantMsg.createdAt).toBe('2025-10-27T10:00:00.000Z'); + expect(assistantMsg.meta).toEqual({ title: 'Test Agent' }); + + // Check tool message + const toolMsg = result.messages[1]; + expect(toolMsg.role).toBe('tool'); + expect(toolMsg.id).toBe('msg-tool-1'); + expect(toolMsg.content).toBe('Weather is sunny, 25°C'); + expect(toolMsg.tool_call_id).toBe('tool-1'); + expect(toolMsg.plugin).toEqual({ + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{"query":"weather"}', + identifier: 'web-browsing', + }); + expect(toolMsg.pluginState).toEqual({ cached: true }); + + // Check metadata + expect(result.metadata.groupMessagesFlattened).toBe(1); + expect(result.metadata.assistantMessagesCreated).toBe(1); + expect(result.metadata.toolMessagesCreated).toBe(1); + }); + + it('should flatten group message with multiple children', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + children: [ + { + id: 'msg-1', + content: 'First response', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-1', + content: 'Result 1', + error: null, + state: {}, + }, + }, + ], + }, + { + id: 'msg-2', + content: 'Follow-up response', + tools: [ + { + id: 'tool-2', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-2', + content: 'Result 2', + error: null, + state: {}, + }, + }, + ], + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Should create 4 messages: 2 assistants + 2 tools + expect(result.messages).toHaveLength(4); + expect(result.messages[0].role).toBe('assistant'); + expect(result.messages[0].id).toBe('msg-1'); + expect(result.messages[1].role).toBe('tool'); + expect(result.messages[1].id).toBe('msg-tool-1'); + expect(result.messages[2].role).toBe('assistant'); + expect(result.messages[2].id).toBe('msg-2'); + expect(result.messages[3].role).toBe('tool'); + expect(result.messages[3].id).toBe('msg-tool-2'); + + expect(result.metadata.groupMessagesFlattened).toBe(1); + expect(result.metadata.assistantMessagesCreated).toBe(2); + expect(result.metadata.toolMessagesCreated).toBe(2); + }); + + it('should handle child without tool result (still executing)', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + children: [ + { + id: 'msg-1', + content: 'Checking weather', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + // No result - tool is still executing + }, + ], + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Should only create 1 assistant message (no tool message) + expect(result.messages).toHaveLength(1); + expect(result.messages[0].role).toBe('assistant'); + expect(result.messages[0].tools).toHaveLength(1); + + expect(result.metadata.assistantMessagesCreated).toBe(1); + expect(result.metadata.toolMessagesCreated).toBe(0); + }); + + it('should handle tool result with error', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + children: [ + { + id: 'msg-1', + content: '', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-1', + content: '', + error: { message: 'Network timeout' }, + state: {}, + }, + }, + ], + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + expect(result.messages).toHaveLength(2); + const toolMsg = result.messages[1]; + expect(toolMsg.role).toBe('tool'); + expect(toolMsg.pluginError).toEqual({ message: 'Network timeout' }); + }); + }); + + describe('Mixed Messages', () => { + it('should preserve non-group messages and flatten group messages', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-user-1', + role: 'user', + content: 'What is the weather?', + }, + { + id: 'msg-group-1', + role: 'group', + content: '', + children: [ + { + id: 'msg-1', + content: 'Checking...', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-1', + content: 'Sunny', + error: null, + state: {}, + }, + }, + ], + }, + ], + }, + { + id: 'msg-user-2', + role: 'user', + content: 'Thanks!', + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // 1 user + (1 assistant + 1 tool) + 1 user = 4 messages + expect(result.messages).toHaveLength(4); + expect(result.messages[0].role).toBe('user'); + expect(result.messages[0].id).toBe('msg-user-1'); + expect(result.messages[1].role).toBe('assistant'); + expect(result.messages[2].role).toBe('tool'); + expect(result.messages[3].role).toBe('user'); + expect(result.messages[3].id).toBe('msg-user-2'); + }); + }); + + describe('Edge Cases', () => { + it('should handle group message with empty children array', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + children: [], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Empty children means no messages created + expect(result.messages).toHaveLength(0); + }); + + it('should handle group message without children field', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + // No children field + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Should keep the message as-is (though this is invalid data) + expect(result.messages).toHaveLength(1); + expect(result.messages[0].id).toBe('msg-group-1'); + }); + + it('should preserve reasoning field from group message', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + reasoning: { + content: 'Thinking about the query...', + signature: 'sig-123', + }, + children: [ + { + id: 'msg-1', + content: 'Result', + tools: [], + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + expect(result.messages).toHaveLength(1); + expect(result.messages[0].reasoning).toEqual({ + content: 'Thinking about the query...', + signature: 'sig-123', + }); + }); + + it('should preserve parent/thread/group/topic IDs', async () => { + const processor = new GroupMessageFlattenProcessor(); + + const input: any[] = [ + { + id: 'msg-group-1', + role: 'group', + content: '', + parentId: 'parent-1', + threadId: 'thread-1', + groupId: 'group-1', + topicId: 'topic-1', + children: [ + { + id: 'msg-1', + content: '', + tools: [ + { + id: 'tool-1', + type: 'builtin', + apiName: 'search', + arguments: '{}', + identifier: 'web-browsing', + result: { + id: 'msg-tool-1', + content: 'Result', + error: null, + state: {}, + }, + }, + ], + }, + ], + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + const assistantMsg = result.messages[0]; + expect(assistantMsg.parentId).toBe('parent-1'); + expect(assistantMsg.threadId).toBe('thread-1'); + expect(assistantMsg.groupId).toBe('group-1'); + expect(assistantMsg.topicId).toBe('topic-1'); + + const toolMsg = result.messages[1]; + expect(toolMsg.parentId).toBe('parent-1'); + expect(toolMsg.threadId).toBe('thread-1'); + expect(toolMsg.groupId).toBe('group-1'); + expect(toolMsg.topicId).toBe('topic-1'); + }); + }); + + describe('Real-world Test Case', () => { + it('should flatten the provided real-world group message', async () => { + const processor = new GroupMessageFlattenProcessor(); + + // Using the real-world test data provided + const input: any[] = [ + { + id: 'msg_LnIlOyMUnX1ylf', + role: 'group', + content: '', + reasoning: { + content: + '**Checking Hangzhou weather**\n\nIt seems the user is asking to check the weather in Hangzhou...', + }, + createdAt: '2025-10-27T10:47:59.475Z', + updatedAt: '2025-10-27T10:48:10.768Z', + topicId: 'tpc_WQ1wRvxdDpLw', + parentId: 'msg_ekwWzxAKueHkd6', + meta: { + avatar: '🤯', + title: '随便聊聊', + }, + children: [ + { + content: '', + id: 'msg_LnIlOyMUnX1ylf', + performance: { + tps: 29.336734693877553, + ttft: 3844, + duration: 3920, + latency: 7764, + }, + tools: [ + { + id: 'call_kYZG2daTTfnkgNiN6oIR25YK', + type: 'builtin', + apiName: 'search', + arguments: + '{"query":"杭州 天气","searchCategories":["general"],"searchEngines":["google","bing"],"searchTimeRange":"day"}', + identifier: 'lobe-web-browsing', + result: { + content: '...', + error: null, + id: 'msg_DS234ZZMju1NNO', + state: { + query: '杭州 天气', + costTime: 1752, + resultNumbers: 600, + }, + }, + }, + ], + usage: { + inputCacheMissTokens: 2404, + totalTokens: 2519, + cost: 0.000831, + }, + }, + ], + usage: { + totalTokens: 2519, + cost: 0.000831, + }, + }, + ]; + + const context = createContext(input); + const result = await processor.process(context); + + // Should create 2 messages + expect(result.messages).toHaveLength(2); + + // Check assistant message + const assistantMsg = result.messages[0]; + expect(assistantMsg.role).toBe('assistant'); + expect(assistantMsg.id).toBe('msg_LnIlOyMUnX1ylf'); + expect(assistantMsg.tools).toHaveLength(1); + expect(assistantMsg.tools[0].identifier).toBe('lobe-web-browsing'); + expect(assistantMsg.tools[0].apiName).toBe('search'); + expect(assistantMsg.reasoning).toBeDefined(); + expect(assistantMsg.topicId).toBe('tpc_WQ1wRvxdDpLw'); + expect(assistantMsg.parentId).toBe('msg_ekwWzxAKueHkd6'); + + // Check tool message + const toolMsg = result.messages[1]; + expect(toolMsg.role).toBe('tool'); + expect(toolMsg.id).toBe('msg_DS234ZZMju1NNO'); + expect(toolMsg.tool_call_id).toBe('call_kYZG2daTTfnkgNiN6oIR25YK'); + expect(toolMsg.plugin).toBeDefined(); + expect(toolMsg.plugin.identifier).toBe('lobe-web-browsing'); + expect(toolMsg.pluginState).toBeDefined(); + expect(toolMsg.pluginState.query).toBe('杭州 天气'); + }); + }); +}); diff --git a/packages/context-engine/src/processors/index.ts b/packages/context-engine/src/processors/index.ts index 5bff1a456d1..39e384886ee 100644 --- a/packages/context-engine/src/processors/index.ts +++ b/packages/context-engine/src/processors/index.ts @@ -1,4 +1,5 @@ // Transformer processors +export { GroupMessageFlattenProcessor } from './GroupMessageFlatten'; export { HistoryTruncateProcessor } from './HistoryTruncate'; export { InputTemplateProcessor } from './InputTemplate'; export { MessageCleanupProcessor } from './MessageCleanup'; diff --git a/packages/database/src/models/__tests__/message.test.ts b/packages/database/src/models/__tests__/message.test.ts index f016fe5a299..bccfd00d77d 100644 --- a/packages/database/src/models/__tests__/message.test.ts +++ b/packages/database/src/models/__tests__/message.test.ts @@ -932,6 +932,78 @@ describe('MessageModel', () => { }); }); + describe('createNewMessage', () => { + it('should create message and return id with messages list', async () => { + // 调用 createNewMessage 方法 + const result = await messageModel.createNewMessage({ + role: 'user', + content: 'test message', + sessionId: '1', + }); + + // 断言返回结构 + expect(result).toHaveProperty('id'); + expect(result).toHaveProperty('messages'); + expect(result.id).toBeDefined(); + expect(result.messages).toBeInstanceOf(Array); + }); + + it('should return newly created message in messages list', async () => { + const content = 'new test message ' + Date.now(); + + const result = await messageModel.createNewMessage({ + role: 'user', + content, + sessionId: '1', + }); + + // 验证新创建的消息在列表中 + const createdMessage = result.messages.find((m) => m.id === result.id); + expect(createdMessage).toBeDefined(); + expect(createdMessage?.content).toBe(content); + expect(createdMessage?.role).toBe('user'); + }); + + it('should return all messages for the session', async () => { + // 创建多条消息 + await messageModel.create({ role: 'user', content: 'message 1', sessionId: '1' }); + await messageModel.create({ role: 'assistant', content: 'message 2', sessionId: '1' }); + + // 创建第三条消息并获取完整列表 + const result = await messageModel.createNewMessage({ + role: 'user', + content: 'message 3', + sessionId: '1', + }); + + // 验证返回了所有消息 + expect(result.messages.length).toBeGreaterThanOrEqual(3); + }); + + it('should filter messages by topicId if provided', async () => { + const topicId = 'topic-1'; + await serverDB.insert(topics).values({ id: topicId, sessionId: '1', userId }); + + // 创建不同 topic 的消息 + await messageModel.create({ role: 'user', content: 'topic 1 msg', sessionId: '1', topicId }); + await messageModel.create({ role: 'user', content: 'no topic msg', sessionId: '1' }); + + // 创建新消息并指定 topicId + const result = await messageModel.createNewMessage({ + role: 'user', + content: 'new topic msg', + sessionId: '1', + topicId, + }); + + // 验证只返回该 topic 的消息 + expect(result.messages.every((m) => m.topicId === topicId || m.topicId === undefined)).toBe( + true, + ); + expect(result.messages.find((m) => m.content === 'no topic msg')).toBeUndefined(); + }); + }); + describe('updateMessage', () => { it('should update message content', async () => { // 创建测试数据 diff --git a/packages/database/src/models/message.ts b/packages/database/src/models/message.ts index 216b8fc7e7c..ad96107b0b5 100644 --- a/packages/database/src/models/message.ts +++ b/packages/database/src/models/message.ts @@ -6,9 +6,11 @@ import { ChatTranslate, ChatVideoItem, CreateMessageParams, + CreateMessageResult, DBMessageItem, ModelRankItem, NewMessageQueryParams, + QueryMessageParams, UIChatMessage, UpdateMessageParams, UpdateMessageRAGParams, @@ -39,14 +41,6 @@ import { LobeChatDatabase } from '../type'; import { genEndDateWhere, genRangeWhere, genStartDateWhere, genWhere } from '../utils/genWhere'; import { idGenerator } from '../utils/idGenerator'; -export interface QueryMessageParams { - current?: number; - groupId?: string | null; - pageSize?: number; - sessionId?: string | null; - topicId?: string | null; -} - export class MessageModel { private userId: string; private db: LobeChatDatabase; @@ -528,6 +522,54 @@ export class MessageModel { }); }; + /** + * Create a new message and return the complete message list + * + * This method combines message creation and querying into a single operation, + * reducing the need for separate refresh calls and improving performance. + * + * @param params - Message creation parameters + * @param options - Query options for post-processing + * @returns Object containing the created message ID and full message list + * + * @example + * const { id, messages } = await messageModel.createNewMessage({ + * role: 'assistant', + * content: 'Hello', + * tools: [...], + * sessionId: 'session-1', + * }); + * // messages already contains grouped structure, no need to refresh + */ + createNewMessage = async ( + params: CreateMessageParams, + options: { + postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise; + } = {}, + ): Promise => { + // 1. Create the message (reuse existing create method) + const item = await this.create(params); + + // 2. Query all messages for this session/topic + // query() method internally applies groupAssistantMessages transformation + const messages = await this.query( + { + current: 0, + groupId: params.groupId, + pageSize: 9999, + sessionId: params.sessionId, + topicId: params.topicId, // Get all messages + }, + options, + ); + + // 3. Return the result + return { + id: item.id, + messages, + }; + }; + batchCreate = async (newMessages: DBMessageItem[]) => { const messagesToInsert = newMessages.map((m) => { // TODO: need a better way to handle this diff --git a/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts b/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts index aa3da936a31..c8f7c77eb8e 100644 --- a/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts +++ b/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts @@ -760,6 +760,40 @@ describe('LobeOpenAICompatibleFactory', () => { } }); + it('should return InsufficientQuota error when error message contains "Insufficient Balance"', async () => { + const apiError = new OpenAI.APIError( + 400, + { + error: { + message: 'Insufficient Balance: Your account balance is too low', + }, + status: 400, + }, + 'Error message', + {}, + ); + + vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError); + + try { + await instance.chat({ + messages: [{ content: 'Hello', role: 'user' }], + model: 'mistralai/mistral-7b-instruct:free', + temperature: 0, + }); + } catch (e) { + expect(e).toEqual({ + endpoint: defaultBaseURL, + error: { + error: { message: 'Insufficient Balance: Your account balance is too low' }, + status: 400, + }, + errorType: AgentRuntimeErrorType.InsufficientQuota, + provider, + }); + } + }); + it('should return AgentRuntimeError for non-OpenAI errors', async () => { // Arrange const genericError = new Error('Generic Error'); diff --git a/packages/model-runtime/src/core/openaiCompatibleFactory/index.ts b/packages/model-runtime/src/core/openaiCompatibleFactory/index.ts index 83897ad0db4..c83accd39fc 100644 --- a/packages/model-runtime/src/core/openaiCompatibleFactory/index.ts +++ b/packages/model-runtime/src/core/openaiCompatibleFactory/index.ts @@ -702,6 +702,18 @@ export const createOpenAICompatibleRuntime = = an log('error code: %s, message: %s', errorResult.code, errorResult.message); + // Check for "Insufficient Balance" in error message + const errorMessage = errorResult.error?.message || errorResult.message; + if (errorMessage?.includes('Insufficient Balance')) { + log('insufficient balance error detected in message'); + return AgentRuntimeError.chat({ + endpoint: desensitizedEndpoint, + error: errorResult, + errorType: AgentRuntimeErrorType.InsufficientQuota, + provider: this.id, + }); + } + switch (errorResult.code) { case 'insufficient_quota': { log('insufficient quota error'); diff --git a/packages/types/src/message/db/params.ts b/packages/types/src/message/db/params.ts index 907497678ff..7982aaa785b 100644 --- a/packages/types/src/message/db/params.ts +++ b/packages/types/src/message/db/params.ts @@ -7,6 +7,32 @@ import { MessageToolCall, ModelReasoning, } from '../common'; +import { UIChatMessage } from '../ui'; + +export interface QueryMessageParams { + current?: number; + groupId?: string | null; + pageSize?: number; + sessionId?: string | null; + topicId?: string | null; +} + +/** + * Result type for createNewMessage + * Contains both the created message ID and the full message list with grouping applied + */ +export interface CreateMessageResult { + /** + * The ID of the created message + */ + id: string; + + /** + * Complete message list with groupAssistantMessages transformation applied + * This includes the newly created message and all existing messages in the session/topic + */ + messages: UIChatMessage[]; +} export interface NewMessage { agentId?: string | null; diff --git a/packages/types/src/message/ui/params.ts b/packages/types/src/message/ui/params.ts index 1cdc70e9a66..401b576aca6 100644 --- a/packages/types/src/message/ui/params.ts +++ b/packages/types/src/message/ui/params.ts @@ -1,6 +1,8 @@ +/* eslint-disable sort-keys-fix/sort-keys-fix , typescript-sort-keys/interface */ import { UploadFileItem } from '../../files'; import { MessageSemanticSearchChunk } from '../../rag'; import { ChatMessageError } from '../common/base'; +import { ChatPluginPayload } from '../common/tools'; import { UIChatMessage, UIMessageRoleType } from './chat'; export interface CreateMessageParams @@ -20,6 +22,44 @@ export interface CreateMessageParams traceId?: string; } +/** + * Parameters for creating a new message with full message list return + * This type is completely independent from UIChatMessage to ensure clean API contract + */ +export interface CreateNewMessageParams { + // ========== Required fields ========== + role: UIMessageRoleType; + content: string; + sessionId: string; + + // ========== Tool related ========== + tool_call_id?: string; + plugin?: ChatPluginPayload; + + // ========== Grouping ========== + parentId?: string; + groupId?: string; + + // ========== Context ========== + topicId?: string; + threadId?: string; + targetId?: string | null; + + // ========== Model info ========== + model?: string; + provider?: string; + + // ========== Content ========== + files?: string[]; + + // ========== Error handling ========== + error?: ChatMessageError | null; + + // ========== Metadata ========== + traceId?: string; + fileChunks?: MessageSemanticSearchChunk[]; +} + export interface SendMessageParams { /** * create a thread diff --git a/src/features/Conversation/Messages/Assistant/Block.tsx b/src/features/Conversation/Messages/Assistant/Block.tsx deleted file mode 100644 index f5aac09e807..00000000000 --- a/src/features/Conversation/Messages/Assistant/Block.tsx +++ /dev/null @@ -1,63 +0,0 @@ -import { AssistantContentBlock } from '@lobechat/types'; -import { Markdown } from '@lobehub/ui'; -import { ReactNode, memo } from 'react'; -import { Flexbox } from 'react-layout-kit'; - -import { LOADING_FLAT } from '@/const/message'; -import ImageFileListViewer from '@/features/Conversation/Messages/User/ImageFileListViewer'; -import { useChatStore } from '@/store/chat'; -import { chatSelectors } from '@/store/chat/selectors'; - -import { DefaultMessage } from '../Default'; -import Tool from './Tool'; - -interface AssistantBlockProps extends AssistantContentBlock { - editableContent: ReactNode; -} -export const AssistantBlock = memo( - ({ id, tools, content, imageList, ...props }) => { - const editing = useChatStore(chatSelectors.isMessageEditing(id)); - const generating = useChatStore(chatSelectors.isMessageGenerating(id)); - - console.log(id, content, props); - const isToolCallGenerating = generating && (content === LOADING_FLAT || !content) && !!tools; - - const showImageItems = !!imageList && imageList.length > 0; - - if (tools && tools.length > 0) - return ( - - {tools.map((toolCall, index) => ( - - ))} - - ); - - if (editing) - return ( - - ); - - return ( - - {content} - {showImageItems && } - - ); - }, -); diff --git a/src/server/routers/lambda/message.ts b/src/server/routers/lambda/message.ts index 377654d75fc..40d4783f028 100644 --- a/src/server/routers/lambda/message.ts +++ b/src/server/routers/lambda/message.ts @@ -66,6 +66,14 @@ export const messageRouter = router({ return data.id; }), + createNewMessage: messageProcedure + .input(z.object({}).passthrough().partial()) + .mutation(async ({ input, ctx }) => { + return ctx.messageModel.createNewMessage(input as any, { + postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path), + }); + }), + // TODO: it will be removed in V2 getAllMessages: messageProcedure.query(async ({ ctx }): Promise => { return ctx.messageModel.queryAll() as any; diff --git a/src/services/message/_deprecated.ts b/src/services/message/_deprecated.ts index 41828340fb1..02ccf160b3f 100644 --- a/src/services/message/_deprecated.ts +++ b/src/services/message/_deprecated.ts @@ -4,6 +4,7 @@ import { ChatTTS, ChatTranslate, CreateMessageParams, + CreateNewMessageParams, ModelRankItem, UIChatMessage, } from '@lobechat/types'; @@ -22,6 +23,13 @@ export class ClientService implements IMessageService { return id; } + async createNewMessage(data: CreateNewMessageParams) { + const { id } = await MessageModel.create(data as any); + const messages = await this.getMessages(data.sessionId, data.topicId); + + return { id, messages }; + } + // @ts-ignore async batchCreateMessages(messages: UIChatMessage[]) { return MessageModel.batchCreate(messages); diff --git a/src/services/message/client.ts b/src/services/message/client.ts index 2f4b4f13cd0..40daf4cf74c 100644 --- a/src/services/message/client.ts +++ b/src/services/message/client.ts @@ -22,6 +22,18 @@ export class ClientService extends BaseClientService implements IMessageService return id; }; + createNewMessage: IMessageService['createNewMessage'] = async ({ sessionId, ...params }) => { + return await this.messageModel.createNewMessage( + { + ...params, + sessionId: sessionId ? (this.toDbSessionId(sessionId) as string) : '', + }, + { + postProcessUrl: this.postProcessUrl, + }, + ); + }; + batchCreateMessages: IMessageService['batchCreateMessages'] = async (messages) => { return this.messageModel.batchCreate(messages); }; @@ -33,12 +45,7 @@ export class ClientService extends BaseClientService implements IMessageService topicId, }, { - postProcessUrl: async (url, file) => { - const hash = (url as string).replace('client-s3://', ''); - const base64 = await this.getBase64ByFileHash(hash); - - return `data:${file.fileType};base64,${base64}`; - }, + postProcessUrl: this.postProcessUrl, }, ); @@ -53,12 +60,7 @@ export class ClientService extends BaseClientService implements IMessageService topicId, }, { - postProcessUrl: async (url, file) => { - const hash = (url as string).replace('client-s3://', ''); - const base64 = await this.getBase64ByFileHash(hash); - - return `data:${file.fileType};base64,${base64}`; - }, + postProcessUrl: this.postProcessUrl, }, ); @@ -168,6 +170,13 @@ export class ClientService extends BaseClientService implements IMessageService return sessionId === INBOX_SESSION_ID ? undefined : sessionId; }; + private postProcessUrl = async (url: string | null, file: any) => { + const hash = (url as string).replace('client-s3://', ''); + const base64 = await this.getBase64ByFileHash(hash); + + return `data:${file.fileType};base64,${base64}`; + }; + private getBase64ByFileHash = async (hash: string) => { const fileItem = await clientS3Storage.getObject(hash); if (!fileItem) throw new Error('file not found'); diff --git a/src/services/message/server.ts b/src/services/message/server.ts index a16e08e3b21..17c278dd789 100644 --- a/src/services/message/server.ts +++ b/src/services/message/server.ts @@ -14,6 +14,13 @@ export class ServerService implements IMessageService { }); }; + createNewMessage: IMessageService['createNewMessage'] = async ({ sessionId, ...params }) => { + return lambdaClient.message.createNewMessage.mutate({ + ...params, + sessionId: sessionId ? this.toDbSessionId(sessionId) : undefined, + }); + }; + batchCreateMessages: IMessageService['batchCreateMessages'] = async (messages) => { return lambdaClient.message.batchCreateMessages.mutate(messages); }; diff --git a/src/services/message/type.ts b/src/services/message/type.ts index 5aded4fac66..5a08e751964 100644 --- a/src/services/message/type.ts +++ b/src/services/message/type.ts @@ -4,6 +4,7 @@ import { ChatTTS, ChatTranslate, CreateMessageParams, + CreateMessageResult, DBMessageItem, ModelRankItem, UIChatMessage, @@ -16,6 +17,7 @@ import type { HeatmapsProps } from '@lobehub/charts'; export interface IMessageService { createMessage(data: CreateMessageParams): Promise; + createNewMessage(data: CreateMessageParams): Promise; batchCreateMessages(messages: DBMessageItem[]): Promise; getMessages(sessionId: string, topicId?: string, groupId?: string): Promise; diff --git a/src/store/chat/slices/plugin/action.test.ts b/src/store/chat/slices/plugin/action.test.ts index d40c89b8680..216765f1d98 100644 --- a/src/store/chat/slices/plugin/action.test.ts +++ b/src/store/chat/slices/plugin/action.test.ts @@ -1,6 +1,5 @@ import { DEFAULT_INBOX_AVATAR, - LOADING_FLAT, PLUGIN_SCHEMA_API_MD5_PREFIX, PLUGIN_SCHEMA_SEPARATOR, } from '@lobechat/const'; @@ -363,16 +362,50 @@ describe('ChatPluginAction', () => { // Verify that tool messages were created for each tool call expect(internal_createMessageMock).toHaveBeenCalledTimes(4); - expect(internal_createMessageMock).toHaveBeenCalledWith({ - content: LOADING_FLAT, + expect(internal_createMessageMock).toHaveBeenNthCalledWith(1, { + content: '', parentId: assistantId, plugin: message.tools![0], role: 'tool', sessionId: 'session-id', tool_call_id: 'tool1', topicId: 'topic-id', + threadId: undefined, + groupId: undefined, + }); + expect(internal_createMessageMock).toHaveBeenNthCalledWith(2, { + content: '', + parentId: assistantId, + plugin: message.tools![1], + role: 'tool', + sessionId: 'session-id', + tool_call_id: 'tool2', + topicId: 'topic-id', + threadId: undefined, + groupId: undefined, + }); + expect(internal_createMessageMock).toHaveBeenNthCalledWith(3, { + content: '', + parentId: assistantId, + plugin: message.tools![2], + role: 'tool', + sessionId: 'session-id', + tool_call_id: 'tool3', + topicId: 'topic-id', + threadId: undefined, + groupId: undefined, + }); + expect(internal_createMessageMock).toHaveBeenNthCalledWith(4, { + content: '', + parentId: assistantId, + plugin: message.tools![3], + role: 'tool', + sessionId: 'session-id', + tool_call_id: 'tool4', + topicId: 'topic-id', + threadId: undefined, + groupId: undefined, }); - // ... similar assertions for other tool calls // Verify that the appropriate plugin types were invoked expect(invokeStandaloneTypePluginMock).toHaveBeenCalledWith( diff --git a/src/store/chat/slices/plugin/action.ts b/src/store/chat/slices/plugin/action.ts index f6123106fe2..d2c9d8990b1 100644 --- a/src/store/chat/slices/plugin/action.ts +++ b/src/store/chat/slices/plugin/action.ts @@ -1,5 +1,4 @@ /* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */ -import { LOADING_FLAT } from '@lobechat/const'; import { ToolNameResolver } from '@lobechat/context-engine'; import { ChatErrorType, @@ -56,6 +55,9 @@ export interface ChatPluginAction { }) => Promise; summaryPluginContent: (id: string) => Promise; + /** + * @deprecated V1 method + */ triggerToolCalls: ( id: string, params?: { threadId?: string; inPortalThread?: boolean; inSearchWorkflow?: boolean }, @@ -258,7 +260,7 @@ export const chatPlugin: StateCreator< let latestToolId = ''; const messagePools = message.tools.map(async (payload) => { const toolMessage: CreateMessageParams = { - content: LOADING_FLAT, + content: '', parentId: assistantId, plugin: payload, role: 'tool',