⚱

src/llms/llm.ts

@@ -227,7 +227,7 @@ export abstract class BaseChatModel<
   }
 
   protected override async _call(
-    ctx: types.TaskCallContext<TInput, TOutput, types.LLMTaskResponseMetadata>
+    ctx: types.TaskCallContext<TInput, types.LLMTaskResponseMetadata>
   ): Promise<types.ParsedData<TOutput>> {
     const messages = await this.buildMessages(ctx.input, ctx)
 

src/task.ts

@@ -1,5 +1,5 @@
 import pRetry, { FailedAttemptError } from 'p-retry'
-import { ZodRawShape, ZodTypeAny } from 'zod'
+import { ZodRawShape, ZodTypeAny, z } from 'zod'
 
 import * as errors from '@/errors'
 import * as types from '@/types'
@@ -65,7 +65,21 @@ export abstract class BaseTask<
   public async callWithMetadata(
     input?: types.ParsedData<TInput>
   ): Promise<types.TaskResponse<TOutput>> {
-    const ctx: types.TaskCallContext<TInput, TOutput> = {
+    if (this.inputSchema) {
+      const inputSchema =
+        this.inputSchema instanceof z.ZodType
+          ? this.inputSchema
+          : z.object(this.inputSchema)
+
+      const safeInput = inputSchema.safeParse(input)
+      if (!safeInput.success) {
+        throw new Error(`Invalid input: ${safeInput.error.message}`)
+      }
+
+      input = safeInput.data
+    }
+
+    const ctx: types.TaskCallContext<TInput> = {
       input,
       attemptNumber: 0,
       metadata: {}
@@ -97,7 +111,7 @@ export abstract class BaseTask<
   }
 
   protected abstract _call(
-    ctx: types.TaskCallContext<TInput, TOutput>
+    ctx: types.TaskCallContext<TInput>
   ): Promise<types.ParsedData<TOutput>>
 
   // TODO

src/tokenizer.ts

@@ -151,13 +151,13 @@ export function getContextSizeForModel(model: string): number {
   }
 }
 
-export const calculateMaxTokens = async ({
+export async function calculateMaxTokens({
   prompt,
   modelName
 }: {
   prompt: string
   modelName: string
-}) => {
+}) {
   // fallback to approximate calculation if tiktoken is not available
   let numTokens = Math.ceil(prompt.length / 4)
 

src/tools/metaphor.ts

@@ -59,19 +59,14 @@ export class MetaphorSearchTool extends BaseTask<
   }
 
   protected override async _call(
-    input: MetaphorSearchToolInput
-  ): Promise<types.TaskResponse<typeof MetaphorSearchToolOutputSchema>> {
-    // TODO: handle errors gracefully
-    input = this.inputSchema.parse(input)
-
+    ctx: types.TaskCallContext<typeof MetaphorSearchToolInputSchema>
+  ): Promise<MetaphorSearchToolOutput> {
+    // TODO: test required inputs
     const result = await this._metaphorClient.search({
-      query: input.query,
-      numResults: input.numResults
+      query: ctx.input!.query,
+      numResults: ctx.input!.numResults
     })
 
-    return {
-      result,
-      metadata: {}
-    }
+    return result
   }
 }

src/tools/novu.ts

@@ -64,15 +64,8 @@ export class NovuNotificationTool extends BaseTask<
   }
 
   protected override async _call(
-    input: NovuNotificationToolInput
-  ): Promise<types.TaskResponse<typeof NovuNotificationToolOutputSchema>> {
-    // TODO: handle errors gracefully
-    input = this.inputSchema.parse(input)
-
-    const result = await this._novuClient.triggerEvent(input)
-    return {
-      result,
-      metadata: {}
-    }
+    ctx: types.TaskCallContext<typeof NovuNotificationToolInputSchema>
+  ): Promise<NovuNotificationToolOutput> {
+    return this._novuClient.triggerEvent(ctx.input!)
   }
 }

src/types.ts

@@ -146,7 +146,6 @@ export interface TaskResponse<
 
 export interface TaskCallContext<
   TInput extends ZodRawShape | ZodTypeAny = ZodTypeAny,
-  TOutput extends ZodRawShape | ZodTypeAny = z.ZodType<string>,
   TMetadata extends TaskResponseMetadata = TaskResponseMetadata
 > {
   input?: ParsedData<TInput>