google-genai[patch]: Fix removing additionalProperties edge case

libs/langchain-google-genai/src/tests/chat_models.test.ts

@@ -1,6 +1,9 @@
 import { test } from "@jest/globals";
 import type { HarmBlockThreshold, HarmCategory } from "@google/generative-ai";
+import { z } from "zod";
+import { zodToJsonSchema } from "zod-to-json-schema";
 import { ChatGoogleGenerativeAI } from "../chat_models.js";
+import { removeAdditionalProperties } from "../utils/zod_to_genai_parameters.js";
 
 test("Google AI - `temperature` must be in range [0.0,1.0]", async () => {
   expect(
@@ -74,3 +77,54 @@ test("Google AI - `safetySettings` category array must be unique", async () => {
       })
   ).toThrow();
 });
+
+test("removeAdditionalProperties can remove all instances of additionalProperties", async () => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  function extractKeys(obj: Record<string, any>, keys: string[] = []) {
+    for (const key in obj) {
+      if (Object.prototype.hasOwnProperty.call(obj, key)) {
+        keys.push(key);
+        if (typeof obj[key] === "object" && obj[key] !== null) {
+          extractKeys(obj[key], keys);
+        }
+      }
+    }
+    return keys;
+  }
+
+  const idealResponseSchema = z.object({
+    idealResponse: z
+      .string()
+      .optional()
+      .describe("The ideal response to the question"),
+  });
+  const questionSchema = z.object({
+    question: z.string().describe("Question text"),
+    type: z.enum(["singleChoice", "multiChoice"]).describe("Question type"),
+    options: z.array(z.string()).describe("List of possible answers"),
+    correctAnswer: z
+      .string()
+      .optional()
+      .describe("correct answer from the possible answers"),
+    idealResponses: z
+      .array(idealResponseSchema)
+      .describe("Array of ideal responses to the question"),
+  });
+
+  const schema = z.object({
+    questions: z.array(questionSchema).describe("Array of question objects"),
+  });
+
+  const parsedSchemaArr = removeAdditionalProperties(zodToJsonSchema(schema));
+  const arrSchemaKeys = extractKeys(parsedSchemaArr);
+  expect(
+    arrSchemaKeys.find((key) => key === "additionalProperties")
+  ).toBeUndefined();
+  const parsedSchemaObj = removeAdditionalProperties(
+    zodToJsonSchema(questionSchema)
+  );
+  const arrSchemaObj = extractKeys(parsedSchemaObj);
+  expect(
+    arrSchemaObj.find((key) => key === "additionalProperties")
+  ).toBeUndefined();
+});

libs/langchain-google-genai/src/utils/zod_to_genai_parameters.ts

@@ -17,34 +17,34 @@ export interface GenerativeAIJsonSchemaDirty extends GenerativeAIJsonSchema {
   additionalProperties?: boolean;
 }
 
-function removeAdditionalProperties(
-  schema: GenerativeAIJsonSchemaDirty
+export function removeAdditionalProperties(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  obj: Record<string, any>
 ): GenerativeAIJsonSchema {
-  const updatedSchema: GenerativeAIJsonSchemaDirty = { ...schema };
-  if (Object.hasOwn(updatedSchema, "additionalProperties")) {
-    delete updatedSchema.additionalProperties;
-  }
-  if (updatedSchema.properties) {
-    const keys = Object.keys(updatedSchema.properties);
-    removeProperties(updatedSchema.properties, keys, 0);
-  }
-
-  return updatedSchema;
-}
-
-function removeProperties(
-  properties: GenerativeAIJsonSchemaDirty["properties"],
-  keys: string[],
-  index: number
-): void {
-  if (index >= keys.length || !properties) {
-    return;
+  if (typeof obj === "object" && obj !== null) {
+    const newObj = { ...obj };
+
+    if (
+      "additionalProperties" in newObj &&
+      typeof newObj.additionalProperties === "boolean"
+    ) {
+      delete newObj.additionalProperties;
+    }
+
+    for (const key in newObj) {
+      if (key in newObj) {
+        if (Array.isArray(newObj[key])) {
+          newObj[key] = newObj[key].map(removeAdditionalProperties);
+        } else if (typeof newObj[key] === "object" && newObj[key] !== null) {
+          newObj[key] = removeAdditionalProperties(newObj[key]);
+        }
+      }
+    }
+
+    return newObj as GenerativeAIJsonSchema;
   }
 
-  const key = keys[index];
-  // eslint-disable-next-line no-param-reassign
-  properties[key] = removeAdditionalProperties(properties[key]);
-  removeProperties(properties, keys, index + 1);
+  return obj as GenerativeAIJsonSchema;
 }
 
 export function zodToGenerativeAIParameters(
@@ -53,9 +53,7 @@ export function zodToGenerativeAIParameters(
 ): GenerativeAIFunctionDeclarationSchema {
   // GenerativeAI doesn't accept either the $schema or additionalProperties
   // attributes, so we need to explicitly remove them.
-  const jsonSchema = removeAdditionalProperties(
-    zodToJsonSchema(zodObj) as GenerativeAIJsonSchemaDirty
-  );
+  const jsonSchema = removeAdditionalProperties(zodToJsonSchema(zodObj));
   const { $schema, ...rest } = jsonSchema;
 
   return rest as GenerativeAIFunctionDeclarationSchema;