task.ts 3.6 KB

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