refactor: improve HumanTool and ConsoleReader implementation

These changes improve the tool's ability to handle interruptions and follow the framework's standard implementation pattern. Ref: i-am-bee#121 Signed-off-by: Matias Molinas <[email protected]>
matiasmolinas · Jan 5, 2025 · 18ce28b · 18ce28b
1 parent f5e74f1
commit 18ce28b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/examples/helpers/io.ts b/examples/helpers/io.ts
@@ -3,6 +3,7 @@ import { stdin, stdout } from "node:process";
 import picocolors from "picocolors";
 import * as R from "remeda";
 import stripAnsi from "strip-ansi";
+import type { Abortable } from 'node:events';
 
 interface ReadFromConsoleInput {
   fallback?: string;
@@ -35,8 +36,11 @@ export function createConsoleReader({
       process.exit(0);
     },
 
-    async askSingleQuestion(queryMessage: string): Promise<string> {
-      const answer = await rl.question(R.piped(picocolors.cyan, picocolors.bold)(queryMessage));
+    async askSingleQuestion(queryMessage: string, options?: Abortable): Promise<string> {
+      const answer = await rl.question(
+        R.piped(picocolors.cyan, picocolors.bold)(queryMessage),
+        options ?? { signal: undefined }
+      );
       return stripAnsi(answer.trim());
     },
 

diff --git a/examples/tools/experimental/human.ts b/examples/tools/experimental/human.ts
@@ -6,6 +6,7 @@ import {
   ToolInput,
   ToolEmitter,
 } from "bee-agent-framework/tools/base";
+import { RunContext } from "bee-agent-framework/context";
 import { z } from "zod";
 
 export class HumanTool extends Tool<StringToolOutput> {
@@ -66,20 +67,21 @@ public readonly emitter: ToolEmitter<ToolInput<this>, StringToolOutput> =
   }
 
   async _run(
-    input: z.infer<ReturnType<typeof this.inputSchema>>,
-    _options: BaseToolRunOptions,
+    input: ToolInput<this>,
+    _options: Partial<BaseToolRunOptions>,
+    run: RunContext<this>
   ): Promise<StringToolOutput> {
     // Use the shared reader instance provided to the constructor
     this.reader.write("HumanTool", input.message);
-
-    // Use askSingleQuestion instead of prompt to avoid interfering with main loop iterator
-    const userInput = await this.reader.askSingleQuestion("User 👤 : ");
-
+  
+    // Use askSingleQuestion with the signal
+    const userInput = await this.reader.askSingleQuestion("User 👤 : ", { signal: run.signal });
+  
     // Format the output as required
     const formattedOutput = `{
       "clarification": "${userInput.trim()}"
     }`;
-
+  
     return new StringToolOutput(formattedOutput);
   }
 }