feat(agent): make bee runner easier to override

Signed-off-by: Tomas Dvorak <[email protected]>
i-am-bee · Oct 28, 2024 · e7b99a0 · e7b99a0
1 parent 8fce700
commit e7b99a0
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/src/agents/bee/agent.ts b/src/agents/bee/agent.ts
@@ -33,10 +33,10 @@ import {
   BeeRunOutput,
 } from "@/agents/bee/types.js";
 import { GetRunContext } from "@/context.js";
-import { BeeAgentRunner } from "@/agents/bee/runner.js";
 import { BeeAgentError } from "@/agents/bee/errors.js";
 import { BeeIterationToolResult } from "@/agents/bee/parser.js";
 import { assign } from "@/internals/helpers/object.js";
+import { BeeAgentRunner } from "@/agents/bee/runner.js";
 
 export interface BeeInput {
   llm: ChatLLM<ChatLLMOutput>;
@@ -52,6 +52,8 @@ export class BeeAgent extends BaseAgent<BeeRunInput, BeeRunOutput, BeeRunOptions
     creator: this,
   });
 
+  protected runner: typeof BeeAgentRunner = BeeAgentRunner;
+
   constructor(protected readonly input: BeeInput) {
     super();
 
@@ -102,7 +104,7 @@ export class BeeAgent extends BaseAgent<BeeRunInput, BeeRunOutput, BeeRunOptions
     const iterations: BeeAgentRunIteration[] = [];
     const maxIterations = options?.execution?.maxIterations ?? Infinity;
 
-    const runner = await BeeAgentRunner.create(this.input, options, input.prompt);
+    const runner = await this.runner.create(this.input, options, input.prompt);
 
     let finalMessage: BaseMessage | undefined;
     while (!finalMessage) {
@@ -195,6 +197,7 @@ export class BeeAgent extends BaseAgent<BeeRunInput, BeeRunOutput, BeeRunOptions
     return {
       input: this.input,
       emitter: this.emitter,
+      runner: this.runner,
     };
   }
 

diff --git a/src/agents/bee/parser.ts b/src/agents/bee/parser.ts
@@ -18,7 +18,7 @@ import { NonUndefined } from "@/internals/types.js";
 import type { LinePrefixParser } from "@/agents/parsers/linePrefix.js";
 import type { BeeAgentRunner } from "@/agents/bee/runner.js";
 
-type Parser = ReturnType<typeof BeeAgentRunner.createParser>["parser"];
+type Parser = ReturnType<InstanceType<typeof BeeAgentRunner>["createParser"]>["parser"];
 export type BeeIterationResult = LinePrefixParser.inferOutput<Parser>;
 export type BeeIterationResultPartial = LinePrefixParser.inferPartialOutput<Parser>;
 export type BeeIterationToolResult = NonUndefined<BeeIterationResult, "tool_input" | "tool_name">;
diff --git a/src/agents/bee/runner.ts b/src/agents/bee/runner.ts
@@ -161,10 +161,10 @@ export class BeeAgentRunner {
       );
     }
 
-    return new BeeAgentRunner(input, options, memory);
+    return new this(input, options, memory);
   }
 
-  static createParser(tools: AnyTool[]) {
+  protected createParser(tools: AnyTool[]) {
     const parserRegex =
       /Thought:.+\n(?:Final Answer:[\S\s]+|Function Name:.+\nFunction Input: \{.*\}\nFunction Caption:.+\nFunction Output:)?/;
 
@@ -251,7 +251,7 @@ export class BeeAgentRunner {
       executor: async () => {
         await emitter.emit("start", { meta });
 
-        const { parser, parserRegex } = BeeAgentRunner.createParser(this.input.tools);
+        const { parser, parserRegex } = this.createParser(this.input.tools);
         const llmOutput = await this.input.llm
           .generate(this.memory.messages.slice(), {
             signal,