task.ts 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  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. import { SessionLock } from "../session/lock"
  10. import { SessionPrompt } from "../session/prompt"
  11. export const TaskTool = Tool.define("task", async () => {
  12. const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary"))
  13. const description = DESCRIPTION.replace(
  14. "{agents}",
  15. agents
  16. .map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
  17. .join("\n"),
  18. )
  19. return {
  20. description,
  21. parameters: z.object({
  22. description: z.string().describe("A short (3-5 words) description of the task"),
  23. prompt: z.string().describe("The task for the agent to perform"),
  24. subagent_type: z.string().describe("The type of specialized agent to use for this task"),
  25. }),
  26. async execute(params, ctx) {
  27. const agent = await Agent.get(params.subagent_type)
  28. if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
  29. const session = await Session.create({
  30. parentID: ctx.sessionID,
  31. title: params.description + ` (@${agent.name} subagent)`,
  32. })
  33. const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
  34. if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
  35. ctx.metadata({
  36. title: params.description,
  37. metadata: {
  38. sessionId: session.id,
  39. },
  40. })
  41. const messageID = Identifier.ascending("message")
  42. const parts: Record<string, MessageV2.ToolPart> = {}
  43. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  44. if (evt.properties.part.sessionID !== session.id) return
  45. if (evt.properties.part.messageID === messageID) return
  46. if (evt.properties.part.type !== "tool") return
  47. parts[evt.properties.part.id] = evt.properties.part
  48. ctx.metadata({
  49. title: params.description,
  50. metadata: {
  51. summary: Object.values(parts).sort((a, b) => a.id?.localeCompare(b.id)),
  52. sessionId: session.id,
  53. },
  54. })
  55. })
  56. const model = agent.model ?? {
  57. modelID: msg.info.modelID,
  58. providerID: msg.info.providerID,
  59. }
  60. ctx.abort.addEventListener("abort", () => {
  61. SessionLock.abort(session.id)
  62. })
  63. const promptParts = await SessionPrompt.resolvePromptParts(params.prompt)
  64. const result = await SessionPrompt.prompt({
  65. messageID,
  66. sessionID: session.id,
  67. model: {
  68. modelID: model.modelID,
  69. providerID: model.providerID,
  70. },
  71. agent: agent.name,
  72. tools: {
  73. todowrite: false,
  74. todoread: false,
  75. task: false,
  76. ...agent.tools,
  77. },
  78. parts: promptParts,
  79. })
  80. unsub()
  81. let all
  82. all = await Session.messages({ sessionID: session.id })
  83. all = all.filter((x) => x.info.role === "assistant")
  84. all = all.flatMap((msg) => msg.parts.filter((x: any) => x.type === "tool") as MessageV2.ToolPart[])
  85. return {
  86. title: params.description,
  87. metadata: {
  88. summary: all,
  89. sessionId: session.id,
  90. },
  91. output: (result.parts.findLast((x: any) => x.type === "text") as any)?.text ?? "",
  92. }
  93. },
  94. }
  95. })