fix: forward tools parameter to Gemini API in GoogleLLM (#4380) (#4386)

This commit is contained in:
dhilip_binny
2026-03-19 02:18:26 +08:00
committed by GitHub
parent a172de9c22
commit 4ffe1eaa4e
2 changed files with 227 additions and 15 deletions
+40 -15
View File
@@ -16,25 +16,50 @@ export class GoogleLLM implements LLM {
responseFormat?: { type: string },
tools?: any[],
): Promise<string | LLMResponse> {
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<string, any> = {};
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```$/, "");
+187
View File
@@ -0,0 +1,187 @@
/// <reference types="jest" />
/**
* 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");
});
});