From b55c51e0049f6fb8dbf4f39b3405e80de905ab22 Mon Sep 17 00:00:00 2001 From: Hrushikesh Yadav <136978914+HrushiYadav@users.noreply.github.com> Date: Mon, 15 Jun 2026 17:28:45 +0530 Subject: [PATCH] fix(anthropic): tool_choice format and tool response parsing (#5537) Co-authored-by: kartik-mem0 --- mem0-ts/src/oss/src/config/manager.ts | 10 + mem0-ts/src/oss/src/llms/anthropic.ts | 60 ++++- mem0-ts/src/oss/src/types/index.ts | 6 + mem0-ts/src/oss/tests/anthropic-llm.test.ts | 263 ++++++++++++++++++++ mem0/llms/anthropic.py | 17 +- 5 files changed, 345 insertions(+), 11 deletions(-) create mode 100644 mem0-ts/src/oss/tests/anthropic-llm.test.ts diff --git a/mem0-ts/src/oss/src/config/manager.ts b/mem0-ts/src/oss/src/config/manager.ts index 27903e707..19dabb083 100644 --- a/mem0-ts/src/oss/src/config/manager.ts +++ b/mem0-ts/src/oss/src/config/manager.ts @@ -112,6 +112,13 @@ export class ConfigManager { | undefined) ?? userConf?.url ?? defaultConf.baseURL; + const llmRaw = userConf as Record | undefined; + const temperature = + userConf?.temperature ?? + (llmRaw?.temperature as number | undefined); + const topP = userConf?.topP ?? (llmRaw?.top_p as number | undefined); + const maxTokens = + userConf?.maxTokens ?? (llmRaw?.max_tokens as number | undefined); return { baseURL: llmBaseURL, @@ -125,6 +132,9 @@ export class ConfigManager { userConf?.modelProperties !== undefined ? userConf.modelProperties : defaultConf.modelProperties, + temperature, + topP, + maxTokens, }; })(), }, diff --git a/mem0-ts/src/oss/src/llms/anthropic.ts b/mem0-ts/src/oss/src/llms/anthropic.ts index f1b13baf9..b5832b0ab 100644 --- a/mem0-ts/src/oss/src/llms/anthropic.ts +++ b/mem0-ts/src/oss/src/llms/anthropic.ts @@ -5,6 +5,9 @@ import { LLMConfig, Message } from "../types"; export class AnthropicLLM implements LLM { private client: Anthropic; private model: string; + private maxTokens: number; + private temperature?: number; + private topP?: number; constructor(config: LLMConfig) { const apiKey = config.apiKey || process.env.ANTHROPIC_API_KEY; @@ -12,18 +15,24 @@ export class AnthropicLLM implements LLM { throw new Error("Anthropic API key is required"); } this.client = new Anthropic({ apiKey }); - this.model = config.model || "claude-3-sonnet-20240229"; + this.model = config.model || "claude-sonnet-4-6"; + // Defaults mirror the Python provider's AnthropicConfig + // (max_tokens=2000, temperature=0.1, top_p omitted). + this.maxTokens = config.maxTokens ?? 2000; + this.temperature = config.temperature ?? 0.1; + this.topP = config.topP; } async generateResponse( messages: Message[], responseFormat?: { type: string }, - ): Promise { + tools?: any[], + ): Promise { // Extract system message if present const systemMessage = messages.find((msg) => msg.role === "system"); const otherMessages = messages.filter((msg) => msg.role !== "system"); - const response = await this.client.messages.create({ + const params: Anthropic.MessageCreateParamsNonStreaming = { model: this.model, messages: otherMessages.map((msg) => ({ role: msg.role as "user" | "assistant", @@ -36,8 +45,41 @@ export class AnthropicLLM implements LLM { typeof systemMessage?.content === "string" ? systemMessage.content : undefined, - max_tokens: 4096, - }); + max_tokens: this.maxTokens, + }; + + // Anthropic rejects requests that include both temperature and top_p; + // prefer temperature, matching the Python provider's _get_common_params. + if (this.temperature !== undefined) { + params.temperature = this.temperature; + } else if (this.topP !== undefined) { + params.top_p = this.topP; + } + + if (tools) { + params.tools = tools; + params.tool_choice = { type: "auto" }; + } + + const response = await this.client.messages.create(params); + + if (tools) { + let content = ""; + const toolCalls: Array<{ name: string; arguments: string }> = []; + + for (const block of response.content) { + if (block.type === "text") { + content = block.text; + } else if (block.type === "tool_use") { + toolCalls.push({ + name: block.name, + arguments: JSON.stringify(block.input), + }); + } + } + + return { content, role: "assistant", toolCalls }; + } const firstBlock = response.content[0]; if (firstBlock.type === "text") { @@ -49,9 +91,9 @@ export class AnthropicLLM implements LLM { async generateChat(messages: Message[]): Promise { const response = await this.generateResponse(messages); - return { - content: response, - role: "assistant", - }; + if (typeof response === "string") { + return { content: response, role: "assistant" }; + } + return response; } } diff --git a/mem0-ts/src/oss/src/types/index.ts b/mem0-ts/src/oss/src/types/index.ts index fd5ec7188..f52bbdb97 100644 --- a/mem0-ts/src/oss/src/types/index.ts +++ b/mem0-ts/src/oss/src/types/index.ts @@ -49,6 +49,9 @@ export interface LLMConfig { model?: string | any; modelProperties?: Record; timeout?: number; + temperature?: number; + topP?: number; + maxTokens?: number; } export interface MemoryConfig { @@ -131,6 +134,9 @@ export const MemoryConfigSchema = z.object({ baseURL: z.string().optional(), url: z.string().optional(), timeout: z.number().optional(), + temperature: z.number().optional(), + topP: z.number().optional(), + maxTokens: z.number().optional(), }), }), historyDbPath: z.string().optional(), diff --git a/mem0-ts/src/oss/tests/anthropic-llm.test.ts b/mem0-ts/src/oss/tests/anthropic-llm.test.ts new file mode 100644 index 000000000..c279037d6 --- /dev/null +++ b/mem0-ts/src/oss/tests/anthropic-llm.test.ts @@ -0,0 +1,263 @@ +/// +/** + * Anthropic LLM — unit tests (mocked @anthropic-ai/sdk). + */ + +const mockCreate = jest.fn(); + +jest.mock("@anthropic-ai/sdk", () => { + return jest.fn().mockImplementation(() => ({ + messages: { create: mockCreate }, + })); +}); + +import { AnthropicLLM } from "../src/llms/anthropic"; + +describe("AnthropicLLM (unit)", () => { + beforeEach(() => mockCreate.mockClear()); + + it("returns text when no tools are provided and model returns a text block", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: '{"facts": ["fact1"]}' }], + }); + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse([ + { role: "user", content: "Hello" }, + ]); + + expect(mockCreate).toHaveBeenCalledTimes(1); + expect(result).toBe('{"facts": ["fact1"]}'); + + // No tools → tool_choice must NOT be forwarded + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.tool_choice).toBeUndefined(); + }); + + // Bug #1 regression: bare string "auto" must NOT be sent; object form required + it("forwards tool_choice as { type: 'auto' } (not bare string) when tools are provided", async () => { + mockCreate.mockResolvedValueOnce({ + content: [ + { + type: "tool_use", + id: "toolu_1", + name: "add_graph_memory", + input: { source: "Alice", destination: "Bob" }, + }, + ], + }); + + const tools = [ + { + name: "add_graph_memory", + description: "Add a graph memory", + input_schema: { type: "object", properties: {} }, + }, + ]; + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + await llm.generateResponse( + [{ role: "user", content: "Alice knows Bob" }], + undefined, + tools, + ); + + const callArgs = mockCreate.mock.calls[0][0]; + // Must be object form, not a bare string + expect(callArgs.tool_choice).toEqual({ type: "auto" }); + expect(callArgs.tool_choice).not.toBe("auto"); + }); + + // Bug #2 regression: must NOT throw on a tool_use block + it("does NOT throw when the model returns a tool_use block", async () => { + mockCreate.mockResolvedValueOnce({ + content: [ + { + type: "tool_use", + id: "toolu_1", + name: "add_graph_memory", + input: { source: "Alice", destination: "Bob" }, + }, + ], + }); + + const tools = [ + { + name: "add_graph_memory", + description: "Add a graph memory", + input_schema: { type: "object", properties: {} }, + }, + ]; + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + await expect( + llm.generateResponse( + [{ role: "user", content: "Alice knows Bob" }], + undefined, + tools, + ), + ).resolves.not.toThrow(); + }); + + it("parses tool_use blocks into toolCalls with JSON-stringified arguments", async () => { + const inputObj = { source: "Alice", destination: "Bob" }; + mockCreate.mockResolvedValueOnce({ + content: [ + { + type: "tool_use", + id: "toolu_1", + name: "add_graph_memory", + input: inputObj, + }, + ], + }); + + const tools = [ + { + name: "add_graph_memory", + description: "Add a graph memory", + input_schema: { type: "object", properties: {} }, + }, + ]; + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Alice knows Bob" }], + undefined, + tools, + ); + + expect(result).toHaveProperty("toolCalls"); + const response = result as { + content: string; + role: string; + toolCalls: Array<{ name: string; arguments: string }>; + }; + expect(response.toolCalls).toHaveLength(1); + expect(response.toolCalls[0].name).toBe("add_graph_memory"); + expect(JSON.parse(response.toolCalls[0].arguments)).toEqual(inputObj); + }); + + it("handles a mixed text + tool_use response", async () => { + const inputObj = { source: "Alice", destination: "Bob" }; + mockCreate.mockResolvedValueOnce({ + content: [ + { type: "text", text: "Calling the tool now." }, + { + type: "tool_use", + id: "toolu_2", + name: "add_graph_memory", + input: inputObj, + }, + ], + }); + + const tools = [ + { + name: "add_graph_memory", + description: "Add a graph memory", + input_schema: { type: "object", properties: {} }, + }, + ]; + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Alice knows Bob" }], + undefined, + tools, + ); + + expect(result).toHaveProperty("toolCalls"); + const response = result as { + content: string; + role: string; + toolCalls: Array<{ name: string; arguments: string }>; + }; + expect(response.content).toBe("Calling the tool now."); + expect(response.role).toBe("assistant"); + expect(response.toolCalls).toHaveLength(1); + expect(response.toolCalls[0].name).toBe("add_graph_memory"); + expect(JSON.parse(response.toolCalls[0].arguments)).toEqual(inputObj); + }); + + it("returns a structured response when tools are provided but the model returns only a text block", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "Just a text response" }], + }); + + const tools = [ + { + name: "noop", + description: "No operation", + input_schema: { type: "object", properties: {} }, + }, + ]; + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Hello" }], + undefined, + tools, + ); + + expect(result).toEqual({ + content: "Just a text response", + role: "assistant", + toolCalls: [], + }); + }); + + // Parity with the Python provider's AnthropicConfig defaults: + // model claude-sonnet-4-6, max_tokens 2000, temperature 0.1, top_p omitted. + it("sends Python-parity defaults (model, max_tokens, temperature)", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "ok" }], + }); + + const llm = new AnthropicLLM({ apiKey: "test-key" }); + await llm.generateResponse([{ role: "user", content: "Hi" }]); + + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.model).toBe("claude-sonnet-4-6"); + expect(callArgs.max_tokens).toBe(2000); + expect(callArgs.temperature).toBe(0.1); + expect(callArgs.top_p).toBeUndefined(); + }); + + it("forwards maxTokens, temperature, and model from config", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "ok" }], + }); + + const llm = new AnthropicLLM({ + apiKey: "test-key", + model: "claude-opus-4-8", + maxTokens: 1024, + temperature: 0.7, + }); + await llm.generateResponse([{ role: "user", content: "Hi" }]); + + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.model).toBe("claude-opus-4-8"); + expect(callArgs.max_tokens).toBe(1024); + expect(callArgs.temperature).toBe(0.7); + }); + + // Anthropic rejects requests with both temperature and top_p set. + it("never sends both temperature and top_p (prefers temperature)", async () => { + mockCreate.mockResolvedValueOnce({ + content: [{ type: "text", text: "ok" }], + }); + + const llm = new AnthropicLLM({ + apiKey: "test-key", + temperature: 0.5, + topP: 0.9, + }); + await llm.generateResponse([{ role: "user", content: "Hi" }]); + + const callArgs = mockCreate.mock.calls[0][0]; + expect(callArgs.temperature).toBe(0.5); + expect(callArgs.top_p).toBeUndefined(); + }); +}); diff --git a/mem0/llms/anthropic.py b/mem0/llms/anthropic.py index 4af33dd8d..e7b2b6810 100644 --- a/mem0/llms/anthropic.py +++ b/mem0/llms/anthropic.py @@ -35,7 +35,7 @@ class AnthropicLLM(LLMBase): super().__init__(config) if not self.config.model: - self.config.model = "claude-3-5-sonnet-20240620" + self.config.model = "claude-sonnet-4-6" api_key = self.config.api_key or os.getenv("ANTHROPIC_API_KEY") self.client = anthropic.Anthropic(api_key=api_key) @@ -106,7 +106,20 @@ class AnthropicLLM(LLMBase): if tools: # TODO: Remove tools if no issues found with new memory addition logic params["tools"] = tools - params["tool_choice"] = tool_choice + params["tool_choice"] = {"type": tool_choice} response = self.client.messages.create(**params) + return self._parse_response(response, tools) + + def _parse_response(self, response, tools): + if tools: + result = {"content": None, "tool_calls": []} + for block in response.content: + if block.type == "text": + result["content"] = block.text + elif block.type == "tool_use": + result["tool_calls"].append( + {"name": block.name, "arguments": block.input} + ) + return result return response.content[0].text