task.ts 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  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 { SessionPrompt } from "../session/prompt"
  10. import { iife } from "@/util/iife"
  11. import { defer } from "@/util/defer"
  12. import { Config } from "../config/config"
  13. export const TaskTool = Tool.define("task", async () => {
  14. const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary"))
  15. const description = DESCRIPTION.replace(
  16. "{agents}",
  17. agents
  18. .map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
  19. .join("\n"),
  20. )
  21. return {
  22. description,
  23. parameters: z.object({
  24. description: z.string().describe("A short (3-5 words) description of the task"),
  25. prompt: z.string().describe("The task for the agent to perform"),
  26. subagent_type: z.string().describe("The type of specialized agent to use for this task"),
  27. session_id: z.string().describe("Existing Task session to continue").optional(),
  28. }),
  29. async execute(params, ctx) {
  30. const agent = await Agent.get(params.subagent_type)
  31. if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
  32. const session = await iife(async () => {
  33. if (params.session_id) {
  34. const found = await Session.get(params.session_id).catch(() => {})
  35. if (found) return found
  36. }
  37. return await Session.create({
  38. parentID: ctx.sessionID,
  39. title: params.description + ` (@${agent.name} subagent)`,
  40. })
  41. })
  42. const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
  43. if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
  44. ctx.metadata({
  45. title: params.description,
  46. metadata: {
  47. sessionId: session.id,
  48. },
  49. })
  50. const messageID = Identifier.ascending("message")
  51. const parts: Record<string, { id: string; tool: string; state: { status: string; title?: string } }> = {}
  52. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  53. if (evt.properties.part.sessionID !== session.id) return
  54. if (evt.properties.part.messageID === messageID) return
  55. if (evt.properties.part.type !== "tool") return
  56. const part = evt.properties.part
  57. parts[part.id] = {
  58. id: part.id,
  59. tool: part.tool,
  60. state: {
  61. status: part.state.status,
  62. title: part.state.status === "completed" ? part.state.title : undefined,
  63. },
  64. }
  65. ctx.metadata({
  66. title: params.description,
  67. metadata: {
  68. summary: Object.values(parts).sort((a, b) => a.id.localeCompare(b.id)),
  69. sessionId: session.id,
  70. },
  71. })
  72. })
  73. const model = agent.model ?? {
  74. modelID: msg.info.modelID,
  75. providerID: msg.info.providerID,
  76. }
  77. function cancel() {
  78. SessionPrompt.cancel(session.id)
  79. }
  80. ctx.abort.addEventListener("abort", cancel)
  81. using _ = defer(() => ctx.abort.removeEventListener("abort", cancel))
  82. const promptParts = await SessionPrompt.resolvePromptParts(params.prompt)
  83. const config = await Config.get()
  84. const result = await SessionPrompt.prompt({
  85. messageID,
  86. sessionID: session.id,
  87. model: {
  88. modelID: model.modelID,
  89. providerID: model.providerID,
  90. },
  91. agent: agent.name,
  92. tools: {
  93. todowrite: false,
  94. todoread: false,
  95. task: false,
  96. ...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])),
  97. ...agent.tools,
  98. },
  99. parts: promptParts,
  100. })
  101. unsub()
  102. const messages = await Session.messages({ sessionID: session.id })
  103. const summary = messages
  104. .filter((x) => x.info.role === "assistant")
  105. .flatMap((msg) => msg.parts.filter((x: any) => x.type === "tool") as MessageV2.ToolPart[])
  106. .map((part) => ({
  107. id: part.id,
  108. tool: part.tool,
  109. state: {
  110. status: part.state.status,
  111. title: part.state.status === "completed" ? part.state.title : undefined,
  112. },
  113. }))
  114. const text = result.parts.findLast((x) => x.type === "text")?.text ?? ""
  115. const output = text + "\n\n" + ["<task_metadata>", `session_id: ${session.id}`, "</task_metadata>"].join("\n")
  116. return {
  117. title: params.description,
  118. metadata: {
  119. summary,
  120. sessionId: session.id,
  121. },
  122. output,
  123. }
  124. },
  125. }
  126. })