feat: supervisor workflow for small models

.env.template

@@ -28,3 +28,7 @@ WATSONX_CHAT_MODEL_OPERATOR="ibm/granite-3-3-8b-instruct"
 BEE_FRAMEWORK_LOG_PRETTY=false
 BEE_FRAMEWORK_LOG_LEVEL="debug"
 BEE_FRAMEWORK_LOG_SINGLE_LINE="false"
+
+# Development only
+DEV_SUPERVISOR_WORKFLOW_ENABLED=true
+DEV_SUPERVISOR_WORKFLOW_DISABLE_EXAMPLES=false

package.json

@@ -59,7 +59,8 @@
     "ui:test:chat": "tsx --inspect --no-warnings examples/ui/chat-monitor/chat-monitor.js",
     "ui:test": "tsx --inspect --no-warnings examples/ui/chat-monitor/input/input.js",
     "ts:check": "tsc --noEmit --project tsconfig.json",
-    "build": "rimraf dist && tsc -b && tsc-alias",
+    "clean": "rimraf dist",
+    "build": "npm run clean && tsc -b tsconfig.build.json && tsc-alias -p tsconfig.build.json",
     "lint": "eslint",
     "lint:fix": "eslint --fix",
     "format": "prettier --check .",
@@ -100,8 +101,9 @@
     "@types/eslint-config-prettier": "^6.11.3",
     "@types/eslint__js": "^8.42.3",
     "@types/node": "^18.15.3",
-    "beeai-framework": "0.1.14",
+    "beeai-framework": "0.1.16",
     "cross-env": "^7.0.3",
+    "date-fns": "^4.1.0",
     "eslint": "^9.10.0",
     "eslint-config-prettier": "^9.1.0",
     "husky": "^9.1.6",
@@ -122,6 +124,6 @@
     "@ai-sdk/openai": "^1.1.14",
     "@eslint/js": "^9.21.0",
     "@i-am-bee/acp-sdk": "0.0.4",
-    "beeai-framework": "0.1.14"
+    "beeai-framework": "0.1.16"
   }
 }

src/agents/agent-factory.ts

@@ -1,46 +1,52 @@
 import { BaseToolsFactory } from "@/base/tools-factory.js";
 import { getChatLLM } from "@/helpers/llm.js";
 import { Switches } from "@/runtime/factory.js";
+import { TaskRunIdValue } from "@/tasks/manager/dto.js";
 import { ReActAgent } from "beeai-framework/agents/react/agent";
+import { AssistantMessage, ToolMessage } from "beeai-framework/backend/message";
 import { TokenMemory } from "beeai-framework/memory/tokenMemory";
-import { UnconstrainedMemory } from "beeai-framework/memory/unconstrainedMemory";
 import { BaseAgentFactory, CreateAgentInput } from "./base/agent-factory.js";
-import { supervisor } from "./index.js";
-import { AssistantMessage, ToolMessage } from "beeai-framework/backend/message";
+import { SupervisorWorkflow } from "./supervisor-workflow/supervisor-workflow.js";
+
+export type AgentUpdateCallback = (key: string, value: string) => void;
 
-export class AgentFactory extends BaseAgentFactory<ReActAgent> {
+export class AgentFactory extends BaseAgentFactory<
+  ReActAgent | SupervisorWorkflow
+> {
   createAgent<TCreateInput extends CreateAgentInput = CreateAgentInput>(
     input: TCreateInput,
     toolsFactory: BaseToolsFactory,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
     switches?: Switches,
   ) {
     const llm = getChatLLM(input.agentKind);
     const generalInstructions = `You are a ${input.agentKind} kind of agent (agentId=${input.agentId}, agentType=${input.agentType}). ${input.instructions}`;
     switch (input.agentKind) {
       case "supervisor": {
-        const tools = toolsFactory.createTools(input.tools);
+        // const tools = toolsFactory.createTools(input.tools);
+        // return new ReActAgent({
+        //   meta: {
+        //     name: input.agentId,
+        //     description: input.description,
+        //   },
+        //   llm,
+        //   memory: new UnconstrainedMemory(),
+        //   tools,
+        //   templates: {
+        //     system: (template) =>
+        //       template.fork((config) => {
+        //         config.defaults.instructions =
+        //           supervisor.SUPERVISOR_INSTRUCTIONS(input.agentId, switches);
+        //       }),
+        //   },
+        //   execution: {
+        //     maxIterations: 100,
+        //     maxRetriesPerStep: 2,
+        //     totalMaxRetries: 10,
+        //   },
+        // });
 
-        return new ReActAgent({
-          meta: {
-            name: input.agentId,
-            description: input.description,
-          },
-          llm,
-          memory: new UnconstrainedMemory(),
-          tools,
-          templates: {
-            system: (template) =>
-              template.fork((config) => {
-                config.defaults.instructions =
-                  supervisor.SUPERVISOR_INSTRUCTIONS(input.agentId, switches);
-              }),
-          },
-          execution: {
-            maxIterations: 100,
-            maxRetriesPerStep: 2,
-            totalMaxRetries: 10,
-          },
-        });
+        return new SupervisorWorkflow(this.logger, llm, input.agentId);
       }
       case "operator":
         return new ReActAgent({
@@ -69,25 +75,29 @@ export class AgentFactory extends BaseAgentFactory<ReActAgent> {
   }
 
   async runAgent(
-    agent: ReActAgent,
+    agent: ReActAgent | SupervisorWorkflow,
     prompt: string,
-    onUpdate: (key: string, value: string) => void,
+    onUpdate: AgentUpdateCallback,
     signal: AbortSignal,
+    taskRunId: TaskRunIdValue,
     addToMemory?: (AssistantMessage | ToolMessage)[],
   ): Promise<string> {
     if (addToMemory) {
       agent.memory.addMany(addToMemory);
     }
+    if (agent instanceof SupervisorWorkflow) {
+      return await agent.run({ prompt, onUpdate, originTaskRunId: taskRunId });
+    } else {
+      const resp = await agent
+        .run({ prompt }, { signal })
+        .observe((emitter) => {
+          emitter.on("update", async ({ update }) => {
+            onUpdate(update.key, update.value);
+          });
+        })
+        .context({ task_run_signal: signal });
 
-    const resp = await agent
-      .run({ prompt }, { signal })
-      .observe((emitter) => {
-        emitter.on("update", async ({ update }) => {
-          onUpdate(update.key, update.value);
-        });
-      })
-      .context({ task_run_signal: signal });
-
-    return resp.result.text;
+      return resp.result.text;
+    }
   }
 }

src/agents/base/agent-factory.ts

@@ -2,6 +2,8 @@ import { BaseToolsFactory } from "@/base/tools-factory.js";
 import { Switches } from "@/runtime/factory.js";
 import { AssistantMessage, ToolMessage } from "beeai-framework/backend/message";
 import { AgentKindEnum } from "../registry/dto.js";
+import { Logger } from "beeai-framework";
+import { TaskRunIdValue } from "@/tasks/manager/dto.js";
 
 export interface CreateAgentInput {
   agentKind: AgentKindEnum;
@@ -13,6 +15,12 @@ export interface CreateAgentInput {
 }
 
 export abstract class BaseAgentFactory<TAgent> {
+  protected logger: Logger;
+  constructor(logger: Logger) {
+    this.logger = logger.child({
+      name: this.constructor.name,
+    });
+  }
   abstract createAgent(
     input: CreateAgentInput,
     toolsFactory: BaseToolsFactory,
@@ -23,6 +31,7 @@ export abstract class BaseAgentFactory<TAgent> {
     prompt: string,
     onUpdate: (key: string, value: string) => void,
     signal: AbortSignal,
+    taskRunId: TaskRunIdValue,
     addToMemory?: (AssistantMessage | ToolMessage)[],
   ): Promise<string>;
 }