From 4ffe1eaa4e14dfa980fdd527e944f843a4568cc1 Mon Sep 17 00:00:00 2001 From: dhilip_binny <33201880+DhilipBinny@users.noreply.github.com> Date: Thu, 19 Mar 2026 02:18:26 +0800 Subject: [PATCH] fix: forward tools parameter to Gemini API in GoogleLLM (#4380) (#4386) --- mem0-ts/src/oss/src/llms/google.ts | 55 +++++-- mem0-ts/src/oss/tests/google-llm.test.ts | 187 +++++++++++++++++++++++ 2 files changed, 227 insertions(+), 15 deletions(-) create mode 100644 mem0-ts/src/oss/tests/google-llm.test.ts diff --git a/mem0-ts/src/oss/src/llms/google.ts b/mem0-ts/src/oss/src/llms/google.ts index 685cef79a..c852baf37 100644 --- a/mem0-ts/src/oss/src/llms/google.ts +++ b/mem0-ts/src/oss/src/llms/google.ts @@ -16,25 +16,50 @@ export class GoogleLLM implements LLM { responseFormat?: { type: string }, tools?: any[], ): Promise { - const completion = await this.google.models.generateContent({ - contents: messages.map((msg) => ({ - parts: [ - { - text: - typeof msg.content === "string" - ? msg.content - : JSON.stringify(msg.content), - }, - ], - role: msg.role === "system" ? "model" : "user", - })), + const contents = messages.map((msg) => ({ + parts: [ + { + text: + typeof msg.content === "string" + ? msg.content + : JSON.stringify(msg.content), + }, + ], + role: msg.role === "system" ? "model" : "user", + })); + // Build config with tools if provided + const config: Record = {}; + if (tools && tools.length > 0) { + config.tools = [ + { + functionDeclarations: tools.map((tool) => ({ + name: tool.function.name, + description: tool.function.description, + parameters: tool.function.parameters, + })), + }, + ]; + } + + const completion = await this.google.models.generateContent({ + contents, model: this.model, - // config: { - // responseSchema: {}, // Add response schema if needed - // }, + config, }); + // Handle function call responses + if (completion.functionCalls && completion.functionCalls.length > 0) { + return { + content: completion.text || "", + role: "assistant", + toolCalls: completion.functionCalls.map((call) => ({ + name: call.name!, + arguments: JSON.stringify(call.args), + })), + }; + } + const text = completion.text ?.replace(/^```json\n/, "") .replace(/\n```$/, ""); diff --git a/mem0-ts/src/oss/tests/google-llm.test.ts b/mem0-ts/src/oss/tests/google-llm.test.ts new file mode 100644 index 000000000..40a3c4e42 --- /dev/null +++ b/mem0-ts/src/oss/tests/google-llm.test.ts @@ -0,0 +1,187 @@ +/// +/** + * Google LLM — unit tests (mocked @google/genai). + * + * Regression tests for #4380: tools parameter was ignored, causing graph + * memory operations to silently fail with Gemini models. + */ + +const mockGenerateContent = jest.fn(); + +jest.mock("@google/genai", () => ({ + GoogleGenAI: jest.fn().mockImplementation(() => ({ + models: { generateContent: mockGenerateContent }, + })), +})); + +import { GoogleLLM } from "../src/llms/google"; + +describe("GoogleLLM (unit)", () => { + beforeEach(() => mockGenerateContent.mockClear()); + + it("returns text response when no tools are provided", async () => { + mockGenerateContent.mockResolvedValueOnce({ + text: '{"facts": ["fact1"]}', + functionCalls: null, + }); + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse([ + { role: "user", content: "Hello" }, + ]); + + expect(mockGenerateContent).toHaveBeenCalledTimes(1); + expect(result).toBe('{"facts": ["fact1"]}'); + + // Verify tools are not in config + const callArgs = mockGenerateContent.mock.calls[0][0]; + expect(callArgs.config.tools).toBeUndefined(); + }); + + it("forwards tools as functionDeclarations to Gemini API", async () => { + mockGenerateContent.mockResolvedValueOnce({ + text: "", + functionCalls: [ + { + name: "extract_entities", + args: { entities: [{ entity: "Alice", entity_type: "person" }] }, + }, + ], + }); + + const tools = [ + { + type: "function", + function: { + name: "extract_entities", + description: "Extract entities from text", + parameters: { + type: "object", + properties: { + entities: { + type: "array", + items: { + type: "object", + properties: { + entity: { type: "string" }, + entity_type: { type: "string" }, + }, + }, + }, + }, + required: ["entities"], + }, + }, + }, + ]; + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Alice is a person" }], + undefined, + tools, + ); + + // Verify functionDeclarations were passed in config + const callArgs = mockGenerateContent.mock.calls[0][0]; + expect(callArgs.config.tools).toBeDefined(); + expect(callArgs.config.tools[0].functionDeclarations).toHaveLength(1); + expect(callArgs.config.tools[0].functionDeclarations[0].name).toBe( + "extract_entities", + ); + + // Verify toolCalls in response + expect(result).toHaveProperty("toolCalls"); + const response = result as { toolCalls: any[] }; + expect(response.toolCalls).toHaveLength(1); + expect(response.toolCalls[0].name).toBe("extract_entities"); + expect(JSON.parse(response.toolCalls[0].arguments)).toEqual({ + entities: [{ entity: "Alice", entity_type: "person" }], + }); + }); + + it("returns text when tools are provided but model returns text", async () => { + mockGenerateContent.mockResolvedValueOnce({ + text: "Just a text response", + functionCalls: null, + }); + + const tools = [ + { + type: "function", + function: { + name: "noop", + description: "No operation", + parameters: { type: "object", properties: {} }, + }, + }, + ]; + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Hello" }], + undefined, + tools, + ); + + // Should return text, not toolCalls + expect(result).toBe("Just a text response"); + }); + + it("strips markdown code fences from text responses", async () => { + mockGenerateContent.mockResolvedValueOnce({ + text: '```json\n{"facts": ["fact1"]}\n```', + functionCalls: null, + }); + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse([ + { role: "user", content: "Extract facts" }, + ]); + + expect(result).toBe('{"facts": ["fact1"]}'); + }); + + it("handles multiple function calls in response", async () => { + mockGenerateContent.mockResolvedValueOnce({ + text: "", + functionCalls: [ + { + name: "add_graph_memory", + args: { source: "Alice", destination: "Bob", relationship: "knows" }, + }, + { + name: "add_graph_memory", + args: { + source: "Bob", + destination: "Charlie", + relationship: "works_with", + }, + }, + ], + }); + + const tools = [ + { + type: "function", + function: { + name: "add_graph_memory", + description: "Add a graph memory", + parameters: { type: "object", properties: {} }, + }, + }, + ]; + + const llm = new GoogleLLM({ apiKey: "test-key" }); + const result = await llm.generateResponse( + [{ role: "user", content: "Alice knows Bob, Bob works with Charlie" }], + undefined, + tools, + ); + + const response = result as { toolCalls: any[] }; + expect(response.toolCalls).toHaveLength(2); + expect(response.toolCalls[0].name).toBe("add_graph_memory"); + expect(response.toolCalls[1].name).toBe("add_graph_memory"); + }); +});