Fix schema input (response_model) for Google AI client (#687)

* Fix schema input (response_model) for Google AI client

* acknowledgement

Author: miguelg719

Diff for: .changeset/short-banks-sit.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+Fixed the schema input for Gemini's response model

Diff for: README.md

@@ -128,7 +128,7 @@ For more information, please see our [Contributing Guide](https://docs.stagehand
 
 ## Acknowledgements
 
-This project heavily relies on [Playwright](https://playwright.dev/) as a resilient backbone to automate the web. It also would not be possible without the awesome techniques and discoveries made by [tarsier](https://github.com/reworkd/tarsier), and [fuji-web](https://github.com/normal-computing/fuji-web).
+This project heavily relies on [Playwright](https://playwright.dev/) as a resilient backbone to automate the web. It also would not be possible without the awesome techniques and discoveries made by [tarsier](https://github.com/reworkd/tarsier), [gemini-zod](https://github.com/jbeoris/gemini-zod), and [fuji-web](https://github.com/normal-computing/fuji-web).
 
 We'd like to thank the following people for their major contributions to Stagehand:
 - [Paul Klein](https://github.com/pkiv)

Diff for: lib/llm/GoogleClient.ts

@@ -14,7 +14,7 @@ import zodToJsonSchema from "zod-to-json-schema";
 import { LogLine } from "../../types/log";
 import { AvailableModel, ClientOptions } from "../../types/model";
 import { LLMCache } from "../cache/LLMCache";
-import { validateZodSchema } from "../utils";
+import { validateZodSchema, toGeminiSchema } from "../utils";
 import {
   ChatCompletionOptions,
   ChatMessage,
@@ -290,25 +290,11 @@ export class GoogleClient extends LLMClient {
       temperature: temperature,
       topP: top_p,
       responseMimeType: response_model ? "application/json" : undefined,
+      responseSchema: response_model
+        ? toGeminiSchema(response_model.schema)
+        : undefined,
     };
 
-    // Handle JSON mode instructions
-    if (response_model) {
-      // Prepend instructions for JSON output if needed (similar to o1 handling)
-      const schemaString = JSON.stringify(
-        zodToJsonSchema(response_model.schema),
-      );
-      formattedMessages.push({
-        role: "user",
-        parts: [
-          {
-            text: `Please respond ONLY with a valid JSON object that strictly adheres to the following JSON schema. Do not include any other text, explanations, or markdown formatting like \`\`\`json ... \`\`\`. Just the JSON object.\n\nSchema:\n${schemaString}`,
-          },
-        ],
-      });
-      formattedMessages.push({ role: "model", parts: [{ text: "{" }] }); // Prime the model
-    }
-
     logger({
       category: "google",
       message: "creating chat completion",

Diff for: lib/utils.ts

@@ -3,6 +3,7 @@ import { z } from "zod";
 import { ObserveResult, Page } from ".";
 import { LogLine } from "../types/log";
 import { TextAnnotation } from "../types/textannotation";
+import { Schema, Type } from "@google/genai";
 
 // This is a heuristic for the width of a character in pixels. It seems to work
 // better than attempting to calculate character widths dynamically, which sometimes
@@ -442,3 +443,123 @@ export function isRunningInBun(): boolean {
     "bun" in process.versions
   );
 }
+
+/*
+ * Helper functions for converting between Gemini and Zod schemas
+ */
+function decorateGeminiSchema(
+  geminiSchema: Schema,
+  zodSchema: z.ZodTypeAny,
+): Schema {
+  if (geminiSchema.nullable === undefined) {
+    geminiSchema.nullable = zodSchema.isOptional();
+  }
+
+  if (zodSchema.description) {
+    geminiSchema.description = zodSchema.description;
+  }
+
+  return geminiSchema;
+}
+
+export function toGeminiSchema(zodSchema: z.ZodTypeAny): Schema {
+  const zodType = getZodType(zodSchema);
+
+  switch (zodType) {
+    case "ZodArray": {
+      return decorateGeminiSchema(
+        {
+          type: Type.ARRAY,
+          items: toGeminiSchema(
+            (zodSchema as z.ZodArray<z.ZodTypeAny>).element,
+          ),
+        },
+        zodSchema,
+      );
+    }
+    case "ZodObject": {
+      const properties: Record<string, Schema> = {};
+      const required: string[] = [];
+
+      Object.entries((zodSchema as z.ZodObject<z.ZodRawShape>).shape).forEach(
+        ([key, value]: [string, z.ZodTypeAny]) => {
+          properties[key] = toGeminiSchema(value);
+          if (getZodType(value) !== "ZodOptional") {
+            required.push(key);
+          }
+        },
+      );
+
+      return decorateGeminiSchema(
+        {
+          type: Type.OBJECT,
+          properties,
+          required: required.length > 0 ? required : undefined,
+        },
+        zodSchema,
+      );
+    }
+    case "ZodString":
+      return decorateGeminiSchema(
+        {
+          type: Type.STRING,
+        },
+        zodSchema,
+      );
+    case "ZodNumber":
+      return decorateGeminiSchema(
+        {
+          type: Type.NUMBER,
+        },
+        zodSchema,
+      );
+    case "ZodBoolean":
+      return decorateGeminiSchema(
+        {
+          type: Type.BOOLEAN,
+        },
+        zodSchema,
+      );
+    case "ZodEnum":
+      return decorateGeminiSchema(
+        {
+          type: Type.STRING,
+          enum: zodSchema._def.values,
+        },
+        zodSchema,
+      );
+    case "ZodDefault":
+    case "ZodNullable":
+    case "ZodOptional": {
+      const innerSchema = toGeminiSchema(zodSchema._def.innerType);
+      return decorateGeminiSchema(
+        {
+          ...innerSchema,
+          nullable: true,
+        },
+        zodSchema,
+      );
+    }
+    case "ZodLiteral":
+      return decorateGeminiSchema(
+        {
+          type: Type.STRING,
+          enum: [zodSchema._def.value],
+        },
+        zodSchema,
+      );
+    default:
+      return decorateGeminiSchema(
+        {
+          type: Type.OBJECT,
+          nullable: true,
+        },
+        zodSchema,
+      );
+  }
+}
+
+// Helper function to check the type of Zod schema
+export function getZodType(schema: z.ZodTypeAny): string {
+  return schema._def.typeName;
+}