Adding zod validation

lib/llm/GoogleClient.ts

@@ -25,8 +25,10 @@ import {
 } from "./LLMClient";
 import {
   CreateChatCompletionResponseError,
+  CreateChatCompletionResponseValidationError,
   StagehandError,
 } from "@/types/stagehandErrors";
+import { validateZodSchemaWithResult } from "@/types/zod";
 
 // Mapping from generic roles to Gemini roles
 const roleMap: { [key in ChatMessage["role"]]: string } = {
@@ -440,7 +442,12 @@ export class GoogleClient extends LLMClient {
           );
         }
 
-        if (!validateZodSchema(response_model.schema, parsedData)) {
+        const validationResult = validateZodSchemaWithResult(
+          response_model.schema,
+          parsedData,
+        );
+
+        if (!validationResult.success) {
           logger({
             category: "google",
             message: "Response failed Zod schema validation",
@@ -453,8 +460,8 @@ export class GoogleClient extends LLMClient {
               retries: retries - 1,
             });
           }
-          throw new CreateChatCompletionResponseError(
-            "Invalid response schema",
+          throw new CreateChatCompletionResponseValidationError(
+            validationResult.error,
           );
         }
 

lib/llm/OpenAIClient.ts

@@ -22,9 +22,10 @@ import {
   LLMResponse,
 } from "./LLMClient";
 import {
-  CreateChatCompletionResponseError,
+  CreateChatCompletionResponseValidationError,
   StagehandError,
 } from "@/types/stagehandErrors";
+import { validateZodSchemaWithResult } from "@/types/zod";
 
 export class OpenAIClient extends LLMClient {
   public type = "openai" as const;
@@ -411,7 +412,12 @@ export class OpenAIClient extends LLMClient {
       const extractedData = response.choices[0].message.content;
       const parsedData = JSON.parse(extractedData);
 
-      if (!validateZodSchema(options.response_model.schema, parsedData)) {
+      const validationResult = validateZodSchemaWithResult(
+        options.response_model.schema,
+        parsedData,
+      );
+
+      if (!validationResult.success) {
         if (retries > 0) {
           // as-casting to account for o1 models not supporting all options
           return this.createChatCompletion({
@@ -421,7 +427,9 @@ export class OpenAIClient extends LLMClient {
           });
         }
 
-        throw new CreateChatCompletionResponseError("Invalid response schema");
+        throw new CreateChatCompletionResponseValidationError(
+          validationResult.error,
+        );
       }
 
       if (this.enableCaching) {

types/stagehandErrors.ts

@@ -1,3 +1,5 @@
+import { ZodValidationError } from "./zod";
+
 export class StagehandError extends Error {
   constructor(message: string) {
     super(message);
@@ -136,6 +138,12 @@ export class CreateChatCompletionResponseError extends StagehandError {
   }
 }
 
+export class CreateChatCompletionResponseValidationError extends StagehandError {
+  constructor(message: ZodValidationError) {
+    super(`ResponseValidationError: ${message.format()}`);
+  }
+}
+
 export class StagehandEvalError extends StagehandError {
   constructor(message: string) {
     super(`StagehandEvalError: ${message}`);

types/zod.ts

@@ -0,0 +1,23 @@
+import { z, ZodError } from "zod";
+
+export interface ZodValidationResult {
+  success: boolean;
+  error?: ZodError;
+}
+
+export function validateZodSchemaWithResult(
+  schema: z.ZodTypeAny,
+  data: unknown
+): ZodValidationResult {
+  try {
+    schema.parse(data);
+    return {
+      success: true,
+    };
+  } catch (error) {
+    return {
+      success: false,
+      error: error instanceof ZodError ? error : new ZodError([]),
+    };
+  }
+}