task.ts 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import { Tool } from "./tool"
  2. import DESCRIPTION from "./task.txt"
  3. import { z } from "zod"
  4. import { Session } from "../session"
  5. import { Bus } from "../bus"
  6. import { MessageV2 } from "../session/message-v2"
  7. import { Identifier } from "../id/id"
  8. import { Agent } from "../agent/agent"
  9. export const TaskTool = Tool.define("task", async () => {
  10. const agents = await Agent.list()
  11. const description = DESCRIPTION.replace("{agents}", agents.map((a) => `- ${a.name}: ${a.description}`).join("\n"))
  12. return {
  13. description,
  14. parameters: z.object({
  15. description: z.string().describe("A short (3-5 words) description of the task"),
  16. prompt: z.string().describe("The task for the agent to perform"),
  17. subagent_type: z.string().describe("The type of specialized agent to use for this task"),
  18. }),
  19. async execute(params, ctx) {
  20. const session = await Session.create(ctx.sessionID)
  21. const msg = await Session.getMessage(ctx.sessionID, ctx.messageID)
  22. if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
  23. const agent = await Agent.get(params.subagent_type)
  24. if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
  25. const messageID = Identifier.ascending("message")
  26. const parts: Record<string, MessageV2.ToolPart> = {}
  27. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  28. if (evt.properties.part.sessionID !== session.id) return
  29. if (evt.properties.part.messageID === messageID) return
  30. if (evt.properties.part.type !== "tool") return
  31. parts[evt.properties.part.id] = evt.properties.part
  32. ctx.metadata({
  33. title: params.description,
  34. metadata: {
  35. summary: Object.values(parts).sort((a, b) => a.id?.localeCompare(b.id)),
  36. },
  37. })
  38. })
  39. const model = agent.model ?? {
  40. modelID: msg.info.modelID,
  41. providerID: msg.info.providerID,
  42. }
  43. ctx.abort.addEventListener("abort", () => {
  44. Session.abort(session.id)
  45. })
  46. const result = await Session.chat({
  47. messageID,
  48. sessionID: session.id,
  49. modelID: model.modelID,
  50. providerID: model.providerID,
  51. mode: msg.info.mode,
  52. system: agent.prompt,
  53. tools: {
  54. ...agent.tools,
  55. task: false,
  56. },
  57. parts: [
  58. {
  59. id: Identifier.ascending("part"),
  60. type: "text",
  61. text: params.prompt,
  62. },
  63. ],
  64. })
  65. unsub()
  66. return {
  67. title: params.description,
  68. metadata: {
  69. summary: result.parts.filter((x) => x.type === "tool"),
  70. },
  71. output: result.parts.findLast((x) => x.type === "text")?.text ?? "",
  72. }
  73. },
  74. }
  75. })