diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/ai-message-sanitizer.test.ts b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/ai-message-sanitizer.test.ts new file mode 100644 index 00000000000..2ea80da9779 --- /dev/null +++ b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/ai-message-sanitizer.test.ts @@ -0,0 +1,114 @@ +/** + * This file is part of the NocoBase (R) project. + * Copyright (c) 2020-2024 NocoBase Co., Ltd. + * Authors: NocoBase Team. + * + * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License. + * For more information, please refer to: https://www.nocobase.com/agreement. + */ + +import { AIMessage } from '@langchain/core/messages'; +import { describe, expect, it, vi } from 'vitest'; +import { sanitizeAdditionalKwargsForToolCalls, sanitizeLangChainAIMessage } from '../ai-employees/tool-call-sanitizer'; +import { convertAIMessage } from '../ai-employees/utils'; + +describe('AI message tool call sanitizer', () => { + const rawToolCall = { + id: 'call_bad', + type: 'function', + function: { + name: 'aiEmployeeWorkflowTaskOutput', + arguments: '{"result":{"reference_reply":":"bad json"}', + }, + }; + + it('should drop malformed raw tool calls when converting LangChain AI messages to stored messages', () => { + const logger = { warn: vi.fn() }; + const aiMessage = new AIMessage({ + id: 'ai_bad', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }); + + const values = convertAIMessage({ + aiEmployee: { + employee: { username: 'assistant' }, + skillSettings: { tools: [] }, + logger, + } as never, + providerName: 'deepseek', + model: 'deepseek-v4-flash', + aiMessage, + }); + + expect(values.metadata.additional_kwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + expect(logger.warn).toHaveBeenCalledWith( + 'Discard malformed raw tool calls from AI message', + expect.objectContaining({ + phase: 'convertAIMessage', + messageId: 'ai_bad', + rawToolCallCount: 1, + parsedToolCallCount: 0, + rawToolCallIds: ['call_bad'], + rawToolCallNames: ['aiEmployeeWorkflowTaskOutput'], + }), + ); + expect(JSON.stringify(logger.warn.mock.calls[0][1])).not.toContain('reference_reply'); + }); + + it('should drop malformed raw tool calls when formatting stored assistant messages', () => { + const { additionalKwargs } = sanitizeAdditionalKwargsForToolCalls( + { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + null, + ); + + expect(additionalKwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + }); + + it('should keep raw tool calls when parsed tool calls are present', () => { + const result = sanitizeAdditionalKwargsForToolCalls( + { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + [{ id: 'call_bad', name: 'aiEmployeeWorkflowTaskOutput', args: {}, type: 'tool_call' }], + ); + + expect(result.changed).toBe(false); + expect(result.additionalKwargs).toEqual({ + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }); + }); + + it('should sanitize LangChain AI messages in place without rebuilding the message instance', () => { + const aiMessage = new AIMessage({ + id: 'ai_bad', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }); + const messageWithRuntimeField = aiMessage as AIMessage & { runtimeField?: string }; + messageWithRuntimeField.runtimeField = 'keep runtime field'; + + const sanitizedMessage = sanitizeLangChainAIMessage(aiMessage); + + expect(sanitizedMessage).toBe(aiMessage); + expect(messageWithRuntimeField.runtimeField).toBe('keep runtime field'); + expect(aiMessage.additional_kwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + }); +}); diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/tool-call-sanitizer-middleware.test.ts b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/tool-call-sanitizer-middleware.test.ts new file mode 100644 index 00000000000..f15da6b2151 --- /dev/null +++ b/packages/plugins/@nocobase/plugin-ai/src/server/__tests__/tool-call-sanitizer-middleware.test.ts @@ -0,0 +1,226 @@ +/** + * This file is part of the NocoBase (R) project. + * Copyright (c) 2020-2024 NocoBase Co., Ltd. + * Authors: NocoBase Team. + * + * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License. + * For more information, please refer to: https://www.nocobase.com/agreement. + */ + +import { AIMessage, BaseMessage, HumanMessage, ToolMessage } from '@langchain/core/messages'; +import { FakeListChatModel } from '@langchain/core/utils/testing'; +import { MemorySaver } from '@langchain/langgraph'; +import { convertMessagesToCompletionsMessageParams } from '@langchain/openai'; +import { createAgent, createMiddleware } from 'langchain'; +import { describe, expect, it, vi } from 'vitest'; +import { toolCallSanitizerMiddleware } from '../ai-employees/middleware'; + +describe('toolCallSanitizerMiddleware', () => { + const rawToolCall = { + id: 'call_bad', + type: 'function', + function: { + name: 'aiEmployeeWorkflowTaskOutput', + arguments: '{"result":{"reference_reply":":"bad json"}', + }, + }; + + type MessagePatchHook = (state: { + messages: BaseMessage[]; + }) => void | { messages?: BaseMessage[] } | Promise; + + const getMessagePatchHook = (hook: unknown): MessagePatchHook => { + if (typeof hook === 'function') { + return hook as MessagePatchHook; + } + if (hook && typeof hook === 'object' && 'hook' in hook) { + const nestedHook = (hook as { hook?: unknown }).hook; + if (typeof nestedHook === 'function') { + return nestedHook as MessagePatchHook; + } + } + throw new Error('Middleware hook is not callable'); + }; + + const getPatchMessages = (result: Awaited>) => { + if (!result || !('messages' in result)) { + return []; + } + return result.messages ?? []; + }; + + it('should remove malformed raw tool calls before checkpoint persistence and next model request', async () => { + const model = new FakeListChatModel({ + responses: [ + new AIMessage({ + id: 'ai_bad', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }), + ], + }); + const logger = { warn: vi.fn() }; + const agent = createAgent({ + model, + tools: [], + middleware: [toolCallSanitizerMiddleware({ logger })], + checkpointer: new MemorySaver(), + }); + const config = { configurable: { thread_id: 'malformed-tool-call' } }; + + await agent.invoke({ messages: [{ role: 'user', content: 'run' }] }, config); + + const state = await agent.getState(config); + const lastMessage = state.values.messages.at(-1) as AIMessage; + + expect(lastMessage.type).toBe('ai'); + expect(lastMessage.tool_calls).toEqual([]); + expect(lastMessage.invalid_tool_calls).toHaveLength(1); + expect(lastMessage.additional_kwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + + const requestMessages = convertMessagesToCompletionsMessageParams({ + messages: state.values.messages, + }); + expect(requestMessages.at(-1)).toMatchObject({ + role: 'assistant', + content: '', + }); + expect(requestMessages.at(-1)).not.toHaveProperty('tool_calls'); + expect(logger.warn).toHaveBeenCalledWith( + 'Discard malformed raw tool calls from AI message', + expect.objectContaining({ + phase: 'afterModel', + messageId: 'ai_bad', + rawToolCallCount: 1, + parsedToolCallCount: 0, + rawToolCallIds: ['call_bad'], + rawToolCallNames: ['aiEmployeeWorkflowTaskOutput'], + }), + ); + expect(JSON.stringify(logger.warn.mock.calls[0][1])).not.toContain('reference_reply'); + }); + + it('should remove malformed raw tool calls from existing state before model requests', async () => { + let capturedMessages: BaseMessage[] = []; + const model = new FakeListChatModel({ + responses: ['done'], + }); + const agent = createAgent({ + model, + tools: [], + middleware: [ + toolCallSanitizerMiddleware(), + createMiddleware({ + name: 'CaptureMessagesMiddleware', + wrapModelCall: (request, handler) => { + capturedMessages = request.messages; + return handler(request); + }, + }), + ], + checkpointer: new MemorySaver(), + }); + + await agent.invoke( + { + messages: [ + new AIMessage({ + id: 'ai_bad_history', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }), + { role: 'user', content: 'continue' }, + ], + }, + { configurable: { thread_id: 'malformed-tool-call-history' } }, + ); + + const requestMessages = convertMessagesToCompletionsMessageParams({ + messages: capturedMessages, + }); + expect(requestMessages[0]).toMatchObject({ + role: 'assistant', + content: '', + }); + expect(requestMessages[0]).not.toHaveProperty('tool_calls'); + expect(requestMessages).toHaveLength(2); + expect(requestMessages[1]).toMatchObject({ + role: 'user', + content: 'continue', + }); + }); + + it('should build a patch without replacing non-AI messages', async () => { + const badAIMessage = new AIMessage({ + id: 'ai_bad_middle', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }); + const humanMessage = new HumanMessage({ id: 'human_1', content: 'continue' }); + const toolMessage = new ToolMessage({ + id: 'tool_1', + content: 'tool result', + tool_call_id: 'call_existing', + }); + const middleware = toolCallSanitizerMiddleware(); + const beforeModel = getMessagePatchHook(middleware.beforeModel); + const messages = [humanMessage, badAIMessage, toolMessage]; + + const result = await beforeModel({ + messages, + }); + const patchMessages = getPatchMessages(result); + + expect(messages).toEqual([humanMessage, badAIMessage, toolMessage]); + expect(patchMessages).toHaveLength(1); + expect(patchMessages[0]).not.toBe(badAIMessage); + expect(AIMessage.isInstance(patchMessages[0])).toBe(true); + expect(patchMessages[0].id).toBe(badAIMessage.id); + expect(badAIMessage.additional_kwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + }); + + it('should sanitize malformed AI messages even when they are not the last message after model', async () => { + const badAIMessage = new AIMessage({ + id: 'ai_bad_not_last', + content: '', + additional_kwargs: { + tool_calls: [rawToolCall], + reasoning_content: 'keep this reasoning', + }, + }); + const middleware = toolCallSanitizerMiddleware(); + const afterModel = getMessagePatchHook(middleware.afterModel); + const messages = [ + new HumanMessage({ id: 'human_1', content: 'start' }), + badAIMessage, + new HumanMessage({ id: 'human_2', content: 'not last' }), + ]; + + const result = await afterModel({ + messages, + }); + const patchMessages = getPatchMessages(result); + + expect(messages[1]).toBe(badAIMessage); + expect(patchMessages).toHaveLength(1); + expect(patchMessages[0]).not.toBe(badAIMessage); + expect(AIMessage.isInstance(patchMessages[0])).toBe(true); + expect(patchMessages[0].id).toBe(badAIMessage.id); + expect(badAIMessage.additional_kwargs).toEqual({ + reasoning_content: 'keep this reasoning', + }); + }); +}); diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts index 4639bd75703..9a5779b8a4d 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/ai-employee.ts @@ -23,6 +23,7 @@ import type { AIEmployee as AIEmployeeType } from '../../collections/ai-employee import { conversationMiddleware, skillToolBindingMiddleware, + toolCallSanitizerMiddleware, toolCallStatusMiddleware, toolInteractionMiddleware, workflowHistoryMiddleware, @@ -39,6 +40,7 @@ import { LLMResult } from '@langchain/core/outputs'; import { Context } from '@nocobase/actions'; import { listAccessibleAIEmployees, serializeEmployeeSummary } from '../../ai/tools/sub-agents/shared'; import { LLMStreamCached } from '../manager/llm-stream-manager'; +import { sanitizeAdditionalKwargsForToolCalls } from './tool-call-sanitizer'; export interface ModelRef { llmService: string; @@ -1288,7 +1290,15 @@ If information is missing, clearly state it in the summary.`; role: 'assistant', content, tool_calls: msg.toolCalls, - additional_kwargs: msg.metadata?.additional_kwargs, + additional_kwargs: sanitizeAdditionalKwargsForToolCalls(msg.metadata?.additional_kwargs, msg.toolCalls, { + onDiscard: (info) => { + this.logger.warn('Discard malformed raw tool calls from AI message', { + phase: 'formatMessages', + messageId: msg.metadata?.id, + ...info, + }); + }, + }).additionalKwargs, }); } @@ -1533,6 +1543,7 @@ If information is missing, clearly state it in the summary.`; toolCallStatusMiddleware(this), ...(inWorkflow ? [workflowHistoryMiddleware(this, this.db)] : []), conversationMiddleware(this, { providerName, llmService, model, messageId, agentThread }), + toolCallSanitizerMiddleware({ logger: this.logger }), ]; } diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/index.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/index.ts index e6a468400ba..ab9c94dcd09 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/index.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/index.ts @@ -9,5 +9,6 @@ export * from './conversation'; export * from './skill-tools'; +export * from './tool-call-sanitizer'; export * from './tools'; export * from './workflow-history'; diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/tool-call-sanitizer.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/tool-call-sanitizer.ts new file mode 100644 index 00000000000..f4e60f1ab0c --- /dev/null +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/middleware/tool-call-sanitizer.ts @@ -0,0 +1,80 @@ +/** + * This file is part of the NocoBase (R) project. + * Copyright (c) 2020-2024 NocoBase Co., Ltd. + * Authors: NocoBase Team. + * + * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License. + * For more information, please refer to: https://www.nocobase.com/agreement. + */ + +import { AIMessage, BaseMessage } from '@langchain/core/messages'; +import { createMiddleware } from 'langchain'; +import { sanitizeLangChainAIMessage } from '../tool-call-sanitizer'; + +type ToolCallSanitizerLogger = { + warn: (message: string, meta?: Record) => void; +}; + +type ToolCallSanitizerMiddlewareOptions = { + logger?: ToolCallSanitizerLogger; +}; + +const buildMessageUpdates = ( + messages: readonly unknown[], + options: ToolCallSanitizerMiddlewareOptions, + phase: 'beforeModel' | 'afterModel', +): BaseMessage[] => { + const updates: BaseMessage[] = []; + + for (const message of messages) { + if (!AIMessage.isInstance(message) || !message.id) { + continue; + } + const sanitizedMessage = sanitizeLangChainAIMessage(message, { + onDiscard: (info) => { + options.logger?.warn('Discard malformed raw tool calls from AI message', { + phase, + messageId: message.id, + invalidToolCallCount: message.invalid_tool_calls?.length ?? 0, + ...info, + }); + }, + }); + if (!sanitizedMessage) { + continue; + } + updates.push( + new AIMessage({ + id: sanitizedMessage.id, + content: sanitizedMessage.content, + name: sanitizedMessage.name, + additional_kwargs: sanitizedMessage.additional_kwargs, + response_metadata: sanitizedMessage.response_metadata, + tool_calls: sanitizedMessage.tool_calls, + invalid_tool_calls: sanitizedMessage.invalid_tool_calls, + usage_metadata: sanitizedMessage.usage_metadata, + }), + ); + } + + return updates; +}; + +export const toolCallSanitizerMiddleware = (options: ToolCallSanitizerMiddlewareOptions = {}) => + createMiddleware({ + name: 'ToolCallSanitizerMiddleware', + beforeModel: (state) => { + const updates = buildMessageUpdates(state.messages ?? [], options, 'beforeModel'); + if (!updates.length) { + return; + } + return { messages: updates }; + }, + afterModel: (state) => { + const updates = buildMessageUpdates(state.messages ?? [], options, 'afterModel'); + if (!updates.length) { + return; + } + return { messages: updates }; + }, + }); diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/tool-call-sanitizer.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/tool-call-sanitizer.ts new file mode 100644 index 00000000000..943fc59bb8c --- /dev/null +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/tool-call-sanitizer.ts @@ -0,0 +1,95 @@ +/** + * This file is part of the NocoBase (R) project. + * Copyright (c) 2020-2024 NocoBase Co., Ltd. + * Authors: NocoBase Team. + * + * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License. + * For more information, please refer to: https://www.nocobase.com/agreement. + */ + +import { AIMessage } from '@langchain/core/messages'; + +type AdditionalKwargs = Record; + +type DiscardedToolCallsInfo = { + rawToolCallCount: number; + parsedToolCallCount: number; + rawToolCallIds: string[]; + rawToolCallNames: string[]; +}; + +type SanitizeAdditionalKwargsOptions = { + onDiscard?: (info: DiscardedToolCallsInfo) => void; +}; + +type SanitizeAdditionalKwargsResult = { + changed: boolean; + additionalKwargs?: AdditionalKwargs; +}; + +const hasRawToolCalls = (additionalKwargs?: AdditionalKwargs) => + Array.isArray(additionalKwargs?.tool_calls) && additionalKwargs.tool_calls.length > 0; + +const hasParsedToolCalls = (toolCalls?: unknown[] | null) => Array.isArray(toolCalls) && toolCalls.length > 0; + +const isRecord = (value: unknown): value is Record => + typeof value === 'object' && value !== null && !Array.isArray(value); + +const collectRawToolCallValues = (toolCalls: unknown[], key: 'id' | 'name') => + toolCalls + .map((toolCall) => { + if (!isRecord(toolCall)) { + return null; + } + if (key === 'name') { + const fn = toolCall.function; + if (isRecord(fn) && typeof fn.name === 'string') { + return fn.name; + } + } + const value = toolCall[key]; + return typeof value === 'string' ? value : null; + }) + .filter((value): value is string => Boolean(value)); + +export const sanitizeAdditionalKwargsForToolCalls = ( + additionalKwargs: AdditionalKwargs | undefined, + parsedToolCalls?: unknown[] | null, + options: SanitizeAdditionalKwargsOptions = {}, +): SanitizeAdditionalKwargsResult => { + if (!hasRawToolCalls(additionalKwargs) || hasParsedToolCalls(parsedToolCalls)) { + return { + changed: false, + additionalKwargs, + }; + } + + const rawToolCalls = additionalKwargs.tool_calls as unknown[]; + options.onDiscard?.({ + rawToolCallCount: rawToolCalls.length, + parsedToolCallCount: parsedToolCalls?.length ?? 0, + rawToolCallIds: collectRawToolCallValues(rawToolCalls, 'id'), + rawToolCallNames: collectRawToolCallValues(rawToolCalls, 'name'), + }); + + const sanitized = { ...additionalKwargs }; + delete sanitized.tool_calls; + + return { + changed: true, + additionalKwargs: Object.keys(sanitized).length > 0 ? sanitized : undefined, + }; +}; + +export const sanitizeLangChainAIMessage = (message: AIMessage, options: SanitizeAdditionalKwargsOptions = {}) => { + const sanitized = sanitizeAdditionalKwargsForToolCalls(message.additional_kwargs, message.tool_calls, options); + + if (!sanitized.changed) { + return null; + } + + message.additional_kwargs = sanitized.additionalKwargs ?? {}; + message.lc_kwargs.additional_kwargs = message.additional_kwargs; + + return message; +}; diff --git a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/utils.ts b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/utils.ts index 366937b6f13..c3a799617b8 100644 --- a/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/utils.ts +++ b/packages/plugins/@nocobase/plugin-ai/src/server/ai-employees/utils.ts @@ -10,6 +10,7 @@ import { AIMessage, HumanMessage, ToolMessage } from 'langchain'; import { AIMessageContent, AIMessageInput } from '../types'; import { AIEmployee } from './ai-employee'; +import { sanitizeAdditionalKwargsForToolCalls } from './tool-call-sanitizer'; export const convertAIMessage = ({ aiEmployee, @@ -84,8 +85,18 @@ export const convertAIMessage = ({ if (aiMessage.response_metadata) { values.metadata.response_metadata = aiMessage.response_metadata; } - if (aiMessage.additional_kwargs) { - values.metadata.additional_kwargs = aiMessage.additional_kwargs; + const additionalKwargs = sanitizeAdditionalKwargsForToolCalls(aiMessage.additional_kwargs, toolCalls, { + onDiscard: (info) => { + aiEmployee.logger?.warn('Discard malformed raw tool calls from AI message', { + phase: 'convertAIMessage', + messageId: aiMessage.id, + invalidToolCallCount: aiMessage.invalid_tool_calls?.length ?? 0, + ...info, + }); + }, + }).additionalKwargs; + if (additionalKwargs) { + values.metadata.additional_kwargs = additionalKwargs; } return values;