@ai-sdk/swarm
Version:
> **Warning** > This is an experimental package. It is not maintained and the API is not stable.
1 lines • 10.9 kB
Source Map (JSON)
{"version":3,"sources":["../src/agent.ts","../src/run-swarm.ts"],"sourcesContent":["import { ToolChoice, LanguageModel, Schema } from 'ai';\nimport { z } from 'zod';\n\ntype Parameters = z.ZodTypeAny | Schema<any>;\ntype inferParameters<PARAMETERS extends Parameters> =\n PARAMETERS extends Schema<any>\n ? PARAMETERS['_type']\n : PARAMETERS extends z.ZodTypeAny\n ? z.infer<PARAMETERS>\n : never;\n\nexport type AgentFunctionTool<\n CONTEXT = any,\n PARAMETERS extends Parameters = any,\n RESULT = any,\n> = {\n type?: undefined | 'function';\n description?: string;\n parameters: PARAMETERS;\n execute: (\n args: inferParameters<PARAMETERS>,\n options: {\n abortSignal?: AbortSignal;\n context: CONTEXT;\n },\n ) => PromiseLike<RESULT>;\n};\n\nexport function functionTool<\n CONTEXT = any,\n PARAMETERS extends Parameters = any,\n>(tool: {\n type?: 'function';\n description?: string;\n parameters: PARAMETERS;\n execute: AgentFunctionTool<CONTEXT, PARAMETERS>['execute'];\n}): AgentFunctionTool<CONTEXT, PARAMETERS> {\n return tool;\n}\n\nexport type AgentHandoverTool<\n CONTEXT = any,\n PARAMETERS extends Parameters = any,\n> = {\n type: 'handover';\n description?: string;\n parameters: PARAMETERS;\n execute: (\n args: inferParameters<PARAMETERS>,\n options: {\n abortSignal?: AbortSignal;\n context: CONTEXT;\n },\n ) => {\n agent: Agent<CONTEXT>;\n context?: CONTEXT;\n };\n};\n\nexport function handoverTool<\n CONTEXT = any,\n PARAMETERS extends Parameters = any,\n>(tool: {\n type: 'handover';\n description?: string;\n parameters: PARAMETERS;\n execute: AgentHandoverTool<CONTEXT, PARAMETERS>['execute'];\n}): AgentHandoverTool<CONTEXT, PARAMETERS> {\n return tool;\n}\n\nexport type AgentTool<CONTEXT = any> =\n | AgentFunctionTool<CONTEXT, any, any>\n | AgentHandoverTool<CONTEXT, any>;\n\n// TODO other settings such as temperature, etc.\nexport class Agent<CONTEXT = any> {\n readonly name: string;\n readonly model: LanguageModel | undefined;\n readonly system: ((context: CONTEXT) => string) | string | undefined;\n readonly tools: Record<string, AgentTool<CONTEXT>> | undefined;\n readonly toolChoice: ToolChoice<any> | undefined;\n\n constructor(options: {\n name: string;\n system?: ((context: CONTEXT) => string) | string | undefined;\n model?: LanguageModel;\n tools?: Record<string, AgentTool>;\n toolChoice?: ToolChoice<any>;\n }) {\n this.name = options.name;\n this.model = options.model;\n this.system = options.system;\n this.tools = options.tools;\n this.toolChoice = options.toolChoice;\n }\n}\n","import {\n CoreAssistantMessage,\n CoreMessage,\n Tool,\n ToolChoice,\n CoreToolMessage,\n FinishReason,\n generateText,\n GenerateTextResult,\n LanguageModel,\n StepResult,\n ToolResultPart,\n} from 'ai';\nimport { Agent, AgentHandoverTool } from './agent';\nimport { z } from 'zod';\n\n// TODO streamSwarm function\nexport async function runSwarm<CONTEXT = any>({\n agent: activeAgent,\n prompt,\n context,\n model,\n maxSteps = 100,\n toolChoice,\n debug = false,\n onStepFinish, // TODO include agent information\n}: {\n agent: Agent;\n prompt: CoreMessage[] | string;\n context?: CONTEXT;\n model: LanguageModel;\n maxSteps?: number;\n toolChoice?: ToolChoice<any>;\n debug?: boolean;\n onStepFinish?: (event: StepResult<any>) => Promise<void> | void;\n}): Promise<{\n text: string;\n responseMessages: CoreMessage[];\n activeAgent: Agent;\n finishReason: FinishReason;\n}> {\n const initialMessages =\n typeof prompt === 'string'\n ? [{ role: 'user' as const, content: prompt }]\n : prompt;\n\n let lastResult: GenerateTextResult<any, any>;\n const responseMessages: Array<CoreMessage> = [];\n\n do {\n lastResult = await generateText({\n model: activeAgent.model ?? model,\n system:\n typeof activeAgent.system === 'function'\n ? activeAgent.system(context)\n : activeAgent.system,\n tools: Object.fromEntries(\n Object.entries(activeAgent.tools ?? {}).map(\n ([name, tool]): [string, Tool] => [\n name,\n tool.type === 'handover'\n ? {\n type: 'function',\n description: tool.description,\n parameters: z.object({}),\n // no execute function\n }\n : {\n type: 'function',\n description: tool.description,\n parameters: tool.parameters,\n execute: (args, { abortSignal }) =>\n tool.execute(args, { context, abortSignal }),\n },\n ],\n ),\n ),\n maxSteps,\n toolChoice: activeAgent.toolChoice ?? toolChoice,\n onStepFinish,\n messages: [...initialMessages, ...responseMessages],\n });\n\n responseMessages.push(...lastResult.response.messages);\n\n // unless there is an agent handover, we are done:\n if (lastResult.finishReason !== 'tool-calls') {\n break;\n }\n\n // the generation stopped with an unhandled tool call\n const { toolCalls, toolResults } = lastResult;\n const toolResultIds = toolResults.map(result => result.toolCallId);\n const unhandledToolCalls = toolCalls.filter(\n toolCall => !toolResultIds.includes(toolCall.toolCallId),\n );\n\n // process handover calls\n const handoverCalls = unhandledToolCalls.filter(\n toolCall => activeAgent.tools?.[toolCall.toolName].type === 'handover',\n );\n\n // take the first handover call (other handover calls are ignored)\n let handoverToolResult: ToolResultPart | undefined = undefined;\n if (handoverCalls.length > 0) {\n const handoverTool = activeAgent.tools?.[\n handoverCalls[0].toolName\n ]! as AgentHandoverTool<CONTEXT, any>;\n\n const result = handoverTool.execute(handoverCalls[0].args, {\n context: context as any,\n });\n\n activeAgent = result.agent;\n context = result.context ?? context; // TODO how to reconcile context?\n\n if (debug) {\n console.log(`\\x1b[36mHanding over to agent ${activeAgent.name}\\x1b[0m`);\n if (result.context != null) {\n console.log(\n `\\x1b[36mUpdated context: ${JSON.stringify(\n result.context,\n null,\n 2,\n )}\\x1b[0m`,\n );\n }\n }\n\n handoverToolResult = {\n type: 'tool-result',\n toolCallId: handoverCalls[0].toolCallId,\n toolName: handoverCalls[0].toolName,\n result: `Handing over to agent ${activeAgent.name}`,\n };\n }\n\n // update last messages\n const toolMessage =\n responseMessages.at(-1)?.role === 'tool'\n ? (responseMessages.at(-1) as CoreToolMessage)\n : undefined;\n const assistantMessage = responseMessages.at(\n toolMessage === undefined ? -1 : -2,\n ) as CoreAssistantMessage;\n\n // add handover tool result\n if (handoverToolResult != null) {\n if (toolMessage == null) {\n responseMessages.push({ role: 'tool', content: [handoverToolResult] });\n } else {\n toolMessage.content.push(handoverToolResult);\n }\n }\n\n // clean out unused tool calls\n if (typeof assistantMessage.content !== 'string') {\n const unusedToolCallIds = handoverCalls\n .filter((call, index) => index > 0)\n .map(call => call.toolCallId);\n\n assistantMessage.content = assistantMessage.content.filter(part => {\n return part.type === 'tool-call'\n ? !unusedToolCallIds.includes(part.toolCallId)\n : true;\n });\n }\n } while (\n responseMessages.filter(message => message.role === 'assistant').length <\n maxSteps\n );\n\n // TODO special finish reasons: done, maxSteps, etc.\n\n return {\n responseMessages,\n activeAgent,\n text: lastResult.text,\n finishReason: lastResult.finishReason,\n };\n}\n"],"mappings":";AA4BO,SAAS,aAGd,MAKyC;AACzC,SAAO;AACT;AAqBO,SAAS,aAGd,MAKyC;AACzC,SAAO;AACT;AAOO,IAAM,QAAN,MAA2B;AAAA,EAOhC,YAAY,SAMT;AACD,SAAK,OAAO,QAAQ;AACpB,SAAK,QAAQ,QAAQ;AACrB,SAAK,SAAS,QAAQ;AACtB,SAAK,QAAQ,QAAQ;AACrB,SAAK,aAAa,QAAQ;AAAA,EAC5B;AACF;;;AChGA;AAAA,EAOE;AAAA,OAKK;AAEP,SAAS,SAAS;AAGlB,eAAsB,SAAwB;AAAA,EAC5C,OAAO;AAAA,EACP;AAAA,EACA;AAAA,EACA;AAAA,EACA,WAAW;AAAA,EACX;AAAA,EACA,QAAQ;AAAA,EACR;AAAA;AACF,GAcG;AAxCH;AAyCE,QAAM,kBACJ,OAAO,WAAW,WACd,CAAC,EAAE,MAAM,QAAiB,SAAS,OAAO,CAAC,IAC3C;AAEN,MAAI;AACJ,QAAM,mBAAuC,CAAC;AAE9C,KAAG;AACD,iBAAa,MAAM,aAAa;AAAA,MAC9B,QAAO,iBAAY,UAAZ,YAAqB;AAAA,MAC5B,QACE,OAAO,YAAY,WAAW,aAC1B,YAAY,OAAO,OAAO,IAC1B,YAAY;AAAA,MAClB,OAAO,OAAO;AAAA,QACZ,OAAO,SAAQ,iBAAY,UAAZ,YAAqB,CAAC,CAAC,EAAE;AAAA,UACtC,CAAC,CAAC,MAAM,IAAI,MAAsB;AAAA,YAChC;AAAA,YACA,KAAK,SAAS,aACV;AAAA,cACE,MAAM;AAAA,cACN,aAAa,KAAK;AAAA,cAClB,YAAY,EAAE,OAAO,CAAC,CAAC;AAAA;AAAA,YAEzB,IACA;AAAA,cACE,MAAM;AAAA,cACN,aAAa,KAAK;AAAA,cAClB,YAAY,KAAK;AAAA,cACjB,SAAS,CAAC,MAAM,EAAE,YAAY,MAC5B,KAAK,QAAQ,MAAM,EAAE,SAAS,YAAY,CAAC;AAAA,YAC/C;AAAA,UACN;AAAA,QACF;AAAA,MACF;AAAA,MACA;AAAA,MACA,aAAY,iBAAY,eAAZ,YAA0B;AAAA,MACtC;AAAA,MACA,UAAU,CAAC,GAAG,iBAAiB,GAAG,gBAAgB;AAAA,IACpD,CAAC;AAED,qBAAiB,KAAK,GAAG,WAAW,SAAS,QAAQ;AAGrD,QAAI,WAAW,iBAAiB,cAAc;AAC5C;AAAA,IACF;AAGA,UAAM,EAAE,WAAW,YAAY,IAAI;AACnC,UAAM,gBAAgB,YAAY,IAAI,YAAU,OAAO,UAAU;AACjE,UAAM,qBAAqB,UAAU;AAAA,MACnC,cAAY,CAAC,cAAc,SAAS,SAAS,UAAU;AAAA,IACzD;AAGA,UAAM,gBAAgB,mBAAmB;AAAA,MACvC,cAAS;AAnGf,YAAAA;AAmGkB,iBAAAA,MAAA,YAAY,UAAZ,gBAAAA,IAAoB,SAAS,UAAU,UAAS;AAAA;AAAA,IAC9D;AAGA,QAAI,qBAAiD;AACrD,QAAI,cAAc,SAAS,GAAG;AAC5B,YAAMC,iBAAe,iBAAY,UAAZ,mBACnB,cAAc,CAAC,EAAE;AAGnB,YAAM,SAASA,cAAa,QAAQ,cAAc,CAAC,EAAE,MAAM;AAAA,QACzD;AAAA,MACF,CAAC;AAED,oBAAc,OAAO;AACrB,iBAAU,YAAO,YAAP,YAAkB;AAE5B,UAAI,OAAO;AACT,gBAAQ,IAAI,iCAAiC,YAAY,IAAI,SAAS;AACtE,YAAI,OAAO,WAAW,MAAM;AAC1B,kBAAQ;AAAA,YACN,4BAA4B,KAAK;AAAA,cAC/B,OAAO;AAAA,cACP;AAAA,cACA;AAAA,YACF,CAAC;AAAA,UACH;AAAA,QACF;AAAA,MACF;AAEA,2BAAqB;AAAA,QACnB,MAAM;AAAA,QACN,YAAY,cAAc,CAAC,EAAE;AAAA,QAC7B,UAAU,cAAc,CAAC,EAAE;AAAA,QAC3B,QAAQ,yBAAyB,YAAY,IAAI;AAAA,MACnD;AAAA,IACF;AAGA,UAAM,gBACJ,sBAAiB,GAAG,EAAE,MAAtB,mBAAyB,UAAS,SAC7B,iBAAiB,GAAG,EAAE,IACvB;AACN,UAAM,mBAAmB,iBAAiB;AAAA,MACxC,gBAAgB,SAAY,KAAK;AAAA,IACnC;AAGA,QAAI,sBAAsB,MAAM;AAC9B,UAAI,eAAe,MAAM;AACvB,yBAAiB,KAAK,EAAE,MAAM,QAAQ,SAAS,CAAC,kBAAkB,EAAE,CAAC;AAAA,MACvE,OAAO;AACL,oBAAY,QAAQ,KAAK,kBAAkB;AAAA,MAC7C;AAAA,IACF;AAGA,QAAI,OAAO,iBAAiB,YAAY,UAAU;AAChD,YAAM,oBAAoB,cACvB,OAAO,CAAC,MAAM,UAAU,QAAQ,CAAC,EACjC,IAAI,UAAQ,KAAK,UAAU;AAE9B,uBAAiB,UAAU,iBAAiB,QAAQ,OAAO,UAAQ;AACjE,eAAO,KAAK,SAAS,cACjB,CAAC,kBAAkB,SAAS,KAAK,UAAU,IAC3C;AAAA,MACN,CAAC;AAAA,IACH;AAAA,EACF,SACE,iBAAiB,OAAO,aAAW,QAAQ,SAAS,WAAW,EAAE,SACjE;AAKF,SAAO;AAAA,IACL;AAAA,IACA;AAAA,IACA,MAAM,WAAW;AAAA,IACjB,cAAc,WAAW;AAAA,EAC3B;AACF;","names":["_a","handoverTool"]}