Add rough Swarm implementation

This commit introduces a new `Swarm` module that includes REPL
functionality and tools for handling functions. Additional changes
include the integration of the `chalk` library for colorful terminal
output. Modifications are made to `package.json` to reflect
these new dependencies and module structures.

package.json

@@ -26,6 +26,10 @@
     "./extract": {
       "types": "./dist/extract/index.d.ts",
       "import": "./dist/extract/index.js"
+    },
+    "./swarm": {
+      "types": "./dist/swarm/index.d.ts",
+      "import": "./dist/swarm/index.js"
     }
   },
   "sideEffects": false,
@@ -68,6 +72,7 @@
     "@dexaai/eslint-config": "^1.3.6",
     "@sentry/node": "^8.34.0",
     "@types/node": "^20.14.11",
+    "chalk": "^5.3.0",
     "dotenv-cli": "^7.4.2",
     "eslint": "^8.57.0",
     "knip": "^5.33.3",

src/swarm/index.ts

@@ -0,0 +1,3 @@
+export { runSwarmRepl } from './repl.js';
+export { Swarm } from './swarm.js';
+export { swarmFunction, swarmHandoff } from './swarm-tools.js';

src/swarm/repl.ts

@@ -0,0 +1,79 @@
+import * as readline from 'node:readline';
+
+import chalk from 'chalk';
+
+import { type Msg, MsgUtil } from '../model/index.js';
+import { Swarm } from './swarm.js';
+import { type Agent } from './types.js';
+
+function prettyPrintMessages(messages: Msg[]): void {
+  for (const message of messages) {
+    // Print tool results
+    if (MsgUtil.isToolResult(message)) {
+      const { content } = message;
+      console.log(`<== ${chalk.green(content)}`);
+    }
+
+    if (message.role !== 'assistant') continue;
+
+    // Print agent name in blue
+    if ('name' in message) {
+      process.stdout.write(`${chalk.blue(message.name || '')}: `);
+    }
+
+    // Print response, if any
+    if (message.content) {
+      console.log(message.content);
+    }
+
+    // Print tool calls in purple, if any
+    const toolCalls = MsgUtil.isToolCall(message) ? message.tool_calls : [];
+    if (toolCalls.length > 1) {
+      console.log();
+    }
+    for (const toolCall of toolCalls) {
+      const { name, arguments: args } = toolCall.function;
+      const argObj = JSON.parse(args);
+      const argStr = JSON.stringify(argObj).replace(/:/g, '=');
+      if (name.startsWith('transfer_')) {
+        console.log(`<> ${chalk.yellow(name)}(${argStr.slice(1, -1)})`);
+      } else {
+        console.log(`${chalk.magenta(name)}(${argStr.slice(1, -1)})`);
+      }
+    }
+  }
+}
+
+export async function runSwarmRepl(
+  startingAgent: Agent,
+  contextVariables: Record<string, any> = {}
+): Promise<void> {
+  const client = new Swarm();
+  console.log('Swarm initialized.');
+
+  let messages: Msg[] = [];
+  let agent = startingAgent;
+
+  const rl = readline.createInterface({
+    input: process.stdin,
+    output: process.stdout,
+  });
+
+  while (true) {
+    const userInput = await new Promise<string>((resolve) => {
+      rl.question(`${chalk.gray('User')}: `, resolve);
+    });
+
+    messages.push({ role: 'user', content: userInput });
+
+    const response = await client.run({
+      agent,
+      messages,
+      ctx: contextVariables,
+    });
+
+    prettyPrintMessages(response.messages);
+    messages = messages.concat(response.messages);
+    agent = response.agent;
+  }
+}

src/swarm/swarm-tools.ts

@@ -0,0 +1,98 @@
+import { z } from 'zod';
+
+import {
+  extractZodObject,
+  type Msg,
+  MsgUtil,
+  zodToJsonSchema,
+} from '../model/index.js';
+import { getErrorMsg } from '../model/utils/errors.js';
+import { cleanString } from '../model/utils/message-util.js';
+import { type Agent, type SwarmFunc } from './types.js';
+
+export function swarmFunction<Schema extends z.ZodObject<any>, Return>(
+  spec: {
+    /** Name of the function. */
+    name: string;
+    /** Description of the function. */
+    description?: string;
+    /** Zod schema for the arguments string. */
+    argsSchema: Schema;
+  },
+  /** Implementation of the function to call with the parsed arguments. */
+  implementation: (params: z.infer<Schema>) => Promise<Return>
+): SwarmFunc {
+  /** Parse the arguments string, optionally reading from a message. */
+  const parseArgs = (input: string | Msg) => {
+    if (typeof input === 'string') {
+      return extractZodObject({ schema: spec.argsSchema, json: input });
+    } else if (MsgUtil.isFuncCall(input)) {
+      const args = input.function_call.arguments;
+      return extractZodObject({ schema: spec.argsSchema, json: args });
+    } else {
+      throw new Error('Invalid input type');
+    }
+  };
+
+  // Call the implementation function with the parsed arguments.
+  const aiFunction = async (input: string | Msg) => {
+    const parsedArgs = parseArgs(input);
+    const result = await implementation(parsedArgs);
+    try {
+      const resultStr =
+        typeof result === 'string' ? result : JSON.stringify(result);
+      return { value: resultStr };
+    } catch (err) {
+      console.error(`Error stringifying function ${spec.name} result:`, err);
+      const errMsg = getErrorMsg(err);
+      return {
+        value: `Error stringifying function ${spec.name} result: ${errMsg}`,
+      };
+    }
+  };
+
+  aiFunction.parseArgs = parseArgs;
+  aiFunction.argsSchema = spec.argsSchema;
+  aiFunction.spec = {
+    name: spec.name,
+    description: cleanString(spec.description ?? ''),
+    parameters: zodToJsonSchema(spec.argsSchema),
+  };
+
+  return aiFunction;
+}
+
+/** This is a simple no-op function that can be used to transfer context to another agent. */
+export function swarmHandoff(args: {
+  agent: Agent;
+  description?: string;
+}): SwarmFunc {
+  const { agent, description } = args;
+  const schema = z.object({});
+
+  /** Parse the arguments string, optionally reading from a message. */
+  const parseArgs = (input: string | Msg) => {
+    if (typeof input === 'string') {
+      return extractZodObject({ schema, json: input });
+    } else if (MsgUtil.isFuncCall(input)) {
+      const args = input.function_call.arguments;
+      return extractZodObject({ schema, json: args });
+    } else {
+      throw new Error(`Invalid input type`);
+    }
+  };
+  const aiFunction = async () => {
+    const value = `Transfered to ${agent.name}. Adopt the role and responsibilities of ${agent.name} and continue the conversation.`;
+    return { value, agent };
+  };
+
+  aiFunction.parseArgs = parseArgs;
+  aiFunction.argsSchema = schema;
+  aiFunction.spec = {
+    name: `transfer_to_${agent.name}`,
+    description: description || '',
+    parameters: zodToJsonSchema(schema),
+  };
+
+  return aiFunction;
+}