Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

google-common[patch]: Zod to Gemini params conversion when the schema contains arrays of items #5553

Merged
37 changes: 36 additions & 1 deletion libs/langchain-google-common/src/tests/utils.test.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
import { expect, test } from "@jest/globals";
import { z } from "zod";
import { zodToGeminiParameters } from "../utils/zod_to_gemini_parameters.js";
Expand All @@ -18,7 +19,7 @@ test("zodToGeminiParameters can convert zod schema to gemini schema", () => {

expect(convertedSchema.type).toBe("object");
expect(convertedSchema.description).toBe("A simple calculator tool");
expect(convertedSchema).not.toContain("additionalProperties");
expect((convertedSchema as any).additionalProperties).toBeUndefined();
expect(convertedSchema.properties).toEqual({
operation: {
type: "string",
Expand All @@ -45,3 +46,37 @@ test("zodToGeminiParameters can convert zod schema to gemini schema", () => {
"childObject",
]);
});

test("zodToGeminiParameters removes additional properties from arrays", () => {
const zodSchema = z
.object({
people: z
.object({
name: z.string().describe("The name of a person"),
})
.array()
.describe("person elements"),
})
.describe("A list of people");

const convertedSchema = zodToGeminiParameters(zodSchema);
expect(convertedSchema.type).toBe("object");
expect(convertedSchema.description).toBe("A list of people");
expect((convertedSchema as any).additionalProperties).toBeUndefined();

const peopleSchema = convertedSchema?.properties?.people;
expect(peopleSchema).not.toBeUndefined();

if (peopleSchema !== undefined) {
expect(peopleSchema.type).toBe("array");
expect((peopleSchema as any).additionalProperties).toBeUndefined();
expect(peopleSchema.description).toBe("person elements");
}

const arrayItemsSchema = peopleSchema?.items;
expect(arrayItemsSchema).not.toBeUndefined();
if (arrayItemsSchema !== undefined) {
expect(arrayItemsSchema.type).toBe("object");
expect((arrayItemsSchema as any).additionalProperties).toBeUndefined();
}
});
1 change: 1 addition & 0 deletions libs/langchain-google-common/src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,7 @@ export type GeminiJsonSchema = Record<string, unknown> & {
};

export interface GeminiJsonSchemaDirty extends GeminiJsonSchema {
items?: GeminiJsonSchemaDirty;
properties?: Record<string, GeminiJsonSchemaDirty>;
additionalProperties?: boolean;
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ function removeAdditionalProperties(
const keys = Object.keys(updatedSchema.properties);
removeProperties(updatedSchema.properties, keys, 0);
}
if (Object.hasOwn(updatedSchema, "items") && updatedSchema.items) {
updatedSchema.items = removeAdditionalProperties(updatedSchema.items);
}

return updatedSchema;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@ import { ChatVertexAI } from "../chat_models.js";

class WeatherTool extends StructuredTool {
schema = z.object({
location: z.string().describe("The name of city to get the weather for."),
locations: z
.array(z.object({ name: z.string() }))
.describe("The name of cities to get the weather for."),
});

description =
Expand All @@ -28,7 +30,7 @@ class WeatherTool extends StructuredTool {

async _call(input: z.infer<typeof this.schema>) {
console.log(`WeatherTool called with input: ${input}`);
return `The weather in ${input.location} is 25°C`;
return `The weather in ${JSON.stringify(input.locations)} is 25°C`;
}
}

Expand Down Expand Up @@ -128,7 +130,7 @@ describe("Google APIKey Chat", () => {

test("Tool call", async () => {
const chat = new ChatVertexAI().bindTools([new WeatherTool()]);
const res = await chat.invoke("What is the weather in SF");
const res = await chat.invoke("What is the weather in SF and LA");
console.log(res);
expect(res.tool_calls?.length).toEqual(1);
expect(res.tool_calls?.[0].args).toEqual(
Expand Down
Loading