feat(agent): improve parser typings

i-am-bee · Sep 26, 2024 · a86d85d · a86d85d
1 parent 2c39c6a
commit a86d85d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git 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
@@ -156,7 +156,6 @@ export class LinePrefixParser<
         if (this.lastNodeKey) {
           const lastNode = this.nodes[this.lastNodeKey];
           if (!lastNode.next.includes(parsedLine.key)) {
-            console.info({ parsedLine, line, x: this.partialState });
             throw new LinePrefixParserError(
               `Transition from '${this.lastNodeKey}' to '${parsedLine.key}' does not exist!`,
             );
@@ -313,9 +312,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;
 }