From 6baed43f12ca3615ed211d26f3a8a0de45d833ba Mon Sep 17 00:00:00 2001
From: Philipp Burckhardt <pburckhardt@outlook.com>
Date: Fri, 16 Jun 2023 17:24:32 -0400
Subject: [PATCH] feat: add before and after hook types

---
 legacy/src/task.ts  | 14 ++++----------
 legacy/src/types.ts |  9 +++++++++
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/legacy/src/task.ts b/legacy/src/task.ts
index d3d8ce73..0054b6c4 100644
--- a/legacy/src/task.ts
+++ b/legacy/src/task.ts
@@ -29,10 +29,8 @@ export abstract class BaseTask<
   protected _timeoutMs?: number
   protected _retryConfig: types.RetryConfig
 
-  private _preHooks: Array<(ctx: types.TaskCallContext<TInput>) => void> = []
-  private _postHooks: Array<
-    (output: TOutput, ctx: types.TaskCallContext<TInput>) => void
-  > = []
+  private _preHooks: Array<types.BeforeCallHook<TInput>> = []
+  private _postHooks: Array<types.AfterCallHook<TInput, TOutput>> = []
 
   constructor(options: types.BaseTaskOptions = {}) {
     this._agentic = options.agentic ?? globalThis.__agentic?.deref()
@@ -79,16 +77,12 @@ export abstract class BaseTask<
     return ''
   }
 
-  public addBeforeCallHook(
-    hook: (ctx: types.TaskCallContext<TInput>) => Promise<void>
-  ): this {
+  public addBeforeCallHook(hook: types.BeforeCallHook<TInput>): this {
     this._preHooks.push(hook)
     return this
   }
 
-  public addAfterCallHook(
-    hook: (output: TOutput, ctx: types.TaskCallContext<TInput>) => Promise<void>
-  ): this {
+  public addAfterCallHook(hook: types.AfterCallHook<TInput, TOutput>): this {
     this._postHooks.push(hook)
     return this
   }
diff --git a/legacy/src/types.ts b/legacy/src/types.ts
index 181bee20..f2541004 100644
--- a/legacy/src/types.ts
+++ b/legacy/src/types.ts
@@ -152,3 +152,12 @@ export declare class CancelablePromise<T> extends Promise<T> {
 }
 
 // export type ProgressFunction = (partialResponse: ChatMessage) => void
+
+export type BeforeCallHook<TInput extends TaskInput = void> = (
+  ctx: TaskCallContext<TInput>
+) => void | Promise<void>
+
+export type AfterCallHook<
+  TInput extends TaskInput = void,
+  TOutput extends TaskOutput = string
+> = (output: TOutput, ctx: TaskCallContext<TInput>) => void | Promise<void>