From 28293a68bad49d00aa1225fe521e6b06303cdc47 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Mon, 9 Dec 2024 15:42:13 +0100
Subject: [PATCH] chore(llms): return driver response type

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 src/llms/drivers/base.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/llms/drivers/base.ts b/src/llms/drivers/base.ts
index 5209dbbc..301c583b 100644
--- a/src/llms/drivers/base.ts
+++ b/src/llms/drivers/base.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { AnySchemaLike, createSchemaValidator, toJsonSchema } from "@/internals/helpers/schema.js";
+import { createSchemaValidator, toJsonSchema } from "@/internals/helpers/schema.js";
 import { GenerateOptions, LLMError } from "@/llms/base.js";
 import { ChatLLM, ChatLLMOutput } from "@/llms/chat.js";
 import { BaseMessage, Role } from "@/llms/primitives/message.js";
@@ -62,7 +62,7 @@ Validation Errors: "{{errors}}"`,
   }
 
   async generate<T = any>(
-    schema: T extends AnySchemaLike ? T : SchemaObject,
+    schema: T extends ZodTypeAny ? T : SchemaObject,
     input: BaseMessage[],
     { maxRetries = 3, options }: GenerateSchemaInput<TGenerateOptions> = {},
   ): Promise<DriverResponse<T>> {