feat: refactor to task

examples/scratch.ts

@@ -16,22 +16,22 @@ async function main() {
   //   .call()
   // console.log(ex0)
 
-  // const ex1 = await $.gpt4(`give me fake data`)
-  //   .output(z.object({ foo: z.string(), bar: z.number() }))
-  //   // .output(z.string())
-  //   // .retry({ attempts: 3 })
-  //   .call()
-  // console.log(ex1)
+  const ex1 = await $.gpt4(`give me fake data`)
+    .output(z.object({ foo: z.string(), bar: z.number() }))
+    // .output(z.string())
+    // .retry({ attempts: 3 })
+    .call()
+  console.log(ex1)
 
-  const getBoolean = $.gpt4(`are you {{mood}}?`)
-    .input(z.object({ mood: z.string() }))
-    .output(z.boolean())
+  // const getBoolean = $.gpt4(`are you {{mood}}?`)
+  //   .input(z.object({ mood: z.string() }))
+  //   .output(z.boolean())
 
-  const results = await Promise.all([
-    getBoolean.call({ mood: 'happy' }),
-    getBoolean.call({ mood: 'sad' })
-  ])
-  console.log(results)
+  // const results = await Promise.all([
+  //   getBoolean.call({ mood: 'happy' }),
+  //   getBoolean.call({ mood: 'sad' })
+  // ])
+  // console.log(results)
 }
 
 main()

src/llm.ts

@@ -1,19 +1,22 @@
 import { ZodRawShape, ZodTypeAny, z } from 'zod'
 
 import * as types from './types'
+import { BaseTaskCallBuilder } from './task'
 
 export abstract class BaseLLMCallBuilder<
   TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
   TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>,
   TModelParams extends Record<string, any> = Record<string, any>
-> {
-  _options: types.BaseLLMOptions<TInput, TOutput, TModelParams>
+> extends BaseTaskCallBuilder<TInput, TOutput> {
+  protected _options: types.BaseLLMOptions<TInput, TOutput, TModelParams>
 
   constructor(options: types.BaseLLMOptions<TInput, TOutput, TModelParams>) {
+    super(options)
+
     this._options = options
   }
 
-  input<U extends ZodRawShape | ZodTypeAny = TInput>(
+  override input<U extends ZodRawShape | ZodTypeAny = TInput>(
     inputSchema: U
   ): BaseLLMCallBuilder<U, TOutput, TModelParams> {
     ;(
@@ -22,7 +25,7 @@ export abstract class BaseLLMCallBuilder<
     return this as unknown as BaseLLMCallBuilder<U, TOutput, TModelParams>
   }
 
-  output<U extends ZodRawShape | ZodTypeAny = TOutput>(
+  override output<U extends ZodRawShape | ZodTypeAny = TOutput>(
     outputSchema: U
   ): BaseLLMCallBuilder<TInput, U, TModelParams> {
     ;(
@@ -46,15 +49,6 @@ export abstract class BaseLLMCallBuilder<
     return this
   }
 
-  retry(retryConfig: types.LLMRetryConfig) {
-    this._options.retryConfig = retryConfig
-    return this
-  }
-
-  abstract call(
-    input?: types.ParsedData<TInput>
-  ): Promise<types.ParsedData<TOutput>>
-
   // TODO
   // abstract stream({
   //   input: TInput,

src/openai.ts

@@ -102,7 +102,7 @@ export class OpenAIChatModelBuilder<
 
         messages.push({
           role: 'system',
-          content: dedent`Output JSON only in the following format:
+          content: dedent`Output JSON only in the following TypeScript format:
           \`\`\`ts
           ${tsTypeString}
           \`\`\``
@@ -162,6 +162,7 @@ export class OpenAIChatModelBuilder<
         if (booleanOutput !== undefined) {
           output = booleanOutput
         } else {
+          // TODO
           throw new Error(`invalid boolean output: ${output}`)
         }
       } else if (outputSchema instanceof z.ZodNumber) {
@@ -172,6 +173,7 @@ export class OpenAIChatModelBuilder<
           : parseFloat(output)
 
         if (isNaN(numberOutput)) {
+          // TODO
           throw new Error(`invalid number output: ${output}`)
         } else {
           output = numberOutput

src/task.ts

@@ -0,0 +1,45 @@
+import { ZodRawShape, ZodTypeAny, z } from 'zod'
+
+import * as types from './types'
+
+export abstract class BaseTaskCallBuilder<
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = z.ZodTypeAny
+> {
+  protected _options: types.BaseTaskOptions<TInput, TOutput>
+
+  constructor(options: types.BaseTaskOptions<TInput, TOutput>) {
+    this._options = options
+  }
+
+  input<U extends ZodRawShape | ZodTypeAny = TInput>(
+    inputSchema: U
+  ): BaseTaskCallBuilder<U, TOutput> {
+    ;(this as unknown as BaseTaskCallBuilder<U, TOutput>)._options.input =
+      inputSchema
+    return this as unknown as BaseTaskCallBuilder<U, TOutput>
+  }
+
+  output<U extends ZodRawShape | ZodTypeAny = TOutput>(
+    outputSchema: U
+  ): BaseTaskCallBuilder<TInput, U> {
+    ;(this as unknown as BaseTaskCallBuilder<TInput, U>)._options.output =
+      outputSchema
+    return this as unknown as BaseTaskCallBuilder<TInput, U>
+  }
+
+  retry(retryConfig: types.RetryConfig) {
+    this._options.retryConfig = retryConfig
+    return this
+  }
+
+  abstract call(
+    input?: types.ParsedData<TInput>
+  ): Promise<types.ParsedData<TOutput>>
+
+  // TODO
+  // abstract stream({
+  //   input: TInput,
+  //   onProgress: types.ProgressFunction
+  // }): Promise<TOutput>
+}

src/types.ts

@@ -24,20 +24,30 @@ export type SafeParsedData<T extends ZodRawShape | ZodTypeAny> =
     ? SafeParseReturnType<ZodObject<T>, ParsedData<T>>
     : never
 
+export interface BaseTaskOptions<
+  TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
+  TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>
+> {
+  input?: TInput
+  output?: TOutput
+
+  timeoutMs?: number
+  retryConfig?: RetryConfig
+
+  // TODO
+  // caching config
+  // logging config
+}
+
 export interface BaseLLMOptions<
   TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
   TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>,
   TModelParams extends Record<string, any> = Record<string, any>
-> {
+> extends BaseTaskOptions<TInput, TOutput> {
   provider?: string
   model?: string
   modelParams?: TModelParams
-  timeoutMs?: number
-
-  input?: TInput
-  output?: TOutput
   examples?: LLMExample[]
-  retryConfig?: LLMRetryConfig
 }
 
 export interface LLMOptions<
@@ -70,7 +80,7 @@ export interface LLMExample {
   output: string
 }
 
-export interface LLMRetryConfig {
+export interface RetryConfig {
   attempts: number
   strategy: string
 }