From a6d9ebe042f61459c590294bb3ccf105d4f5d2c2 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Thu, 26 Sep 2024 23:22:12 +0200
Subject: [PATCH] feat(agent): improve parser typings

---
 src/agents/parsers/field.ts      | 21 +++++++++------------
 src/agents/parsers/linePrefix.ts | 13 ++++++++++++-
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/src/agents/parsers/field.ts b/src/agents/parsers/field.ts
index 3ec9082b..79583f4f 100644
--- a/src/agents/parsers/field.ts
+++ b/src/agents/parsers/field.ts
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { z, ZodSchema } from "zod";
+import { ZodSchema } from "zod";
 import { setProp } from "@/internals/helpers/object.js";
 import { ValueError } from "@/errors.js";
 import { Serializable } from "@/internals/serializable.js";
@@ -44,20 +44,20 @@ export abstract class ParserField<T, TPartial> extends Serializable {
   }
 }
 
-export class ZodParserField<T extends ZodSchema> extends ParserField<z.output<T>, string> {
+export class ZodParserField<T> extends ParserField<T, string> {
   static {
     this.register();
   }
 
-  constructor(protected readonly schema: T) {
+  constructor(protected readonly schema: ZodSchema<T>) {
     super();
   }
 
-  get() {
+  get(): T {
     return this.schema.parse(this.raw);
   }
 
-  getPartial() {
+  getPartial(): string {
     return this.raw;
   }
 
@@ -66,15 +66,12 @@ export class ZodParserField<T extends ZodSchema> extends ParserField<z.output<T>
   }
 }
 
-export class JSONParserField<T extends ZodSchema> extends ParserField<
-  z.output<T>,
-  Partial<z.output<T>>
-> {
+export class JSONParserField<T> extends ParserField<T, Partial<T>> {
   protected stream!: ReturnType<typeof jsonrepairTransform>;
   protected jsonParser!: JSONParser;
-  protected ref!: { value: z.output<T> };
+  protected ref!: { value: Partial<T> };
 
-  constructor(protected readonly input: { schema: T; base: Partial<z.output<T>> }) {
+  constructor(protected readonly input: { schema: ZodSchema<T>; base: Partial<T> }) {
     super();
     if (input.base === undefined) {
       throw new ValueError(`Base must be defined!`);
@@ -145,6 +142,6 @@ export class JSONParserField<T extends ZodSchema> extends ParserField<
 
 // eslint-disable-next-line @typescript-eslint/no-namespace
 export namespace ParserField {
-  export type inferValue<T> = T extends ParserField<infer L, any> ? L : never;
+  export type inferValue<T> = T extends ParserField<infer L, unknown> ? L : never;
   export type inferPartialValue<T> = T extends ParserField<any, infer L> ? L : never;
 }
diff --git a/src/agents/parsers/linePrefix.ts b/src/agents/parsers/linePrefix.ts
index 2456c73e..86d51c8f 100644
--- a/src/agents/parsers/linePrefix.ts
+++ b/src/agents/parsers/linePrefix.ts
@@ -313,9 +313,20 @@ export namespace LinePrefixParser {
   export type infer<
     T extends Record<string, ParserNode<Extract<keyof T, string>, ParserField<any, any>>>,
   > = {
-    [K in StringKey<T>]: T[K]["field"] extends ParserField<infer L, any> ? L : never;
+    [K in keyof T]: ParserField.inferValue<T[K]["field"]>;
+  };
+  export type inferPartial<
+    T extends Record<string, ParserNode<Extract<keyof T, string>, ParserField<any, any>>>,
+  > = {
+    [K in StringKey<T>]: ParserField.inferPartialValue<T[K]["field"]>;
   };
 
+  export type inferCallback<T> =
+    T extends LinePrefixParser<any> ? (T["emitter"] extends Emitter<infer L> ? L : never) : never;
+
   export type inferOutput<T> =
     T extends LinePrefixParser<infer P> ? LinePrefixParser.infer<P> : never;
+
+  export type inferPartialOutput<T> =
+    T extends LinePrefixParser<infer P> ? LinePrefixParser.inferPartial<P> : never;
 }