task.ts 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  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. export const TaskTool = Tool.define({
  9. id: "task",
  10. description: DESCRIPTION,
  11. parameters: z.object({
  12. description: z.string().describe("A short (3-5 words) description of the task"),
  13. prompt: z.string().describe("The task for the agent to perform"),
  14. }),
  15. async execute(params, ctx) {
  16. const session = await Session.create(ctx.sessionID)
  17. const msg = await Session.getMessage(ctx.sessionID, ctx.messageID)
  18. if (msg.role !== "assistant") throw new Error("Not an assistant message")
  19. const messageID = Identifier.ascending("message")
  20. const parts: Record<string, MessageV2.ToolPart> = {}
  21. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  22. if (evt.properties.part.sessionID !== session.id) return
  23. if (evt.properties.part.messageID === messageID) return
  24. if (evt.properties.part.type !== "tool") return
  25. parts[evt.properties.part.id] = evt.properties.part
  26. ctx.metadata({
  27. title: params.description,
  28. metadata: {
  29. summary: Object.values(parts).sort((a, b) => a.id?.localeCompare(b.id)),
  30. },
  31. })
  32. })
  33. ctx.abort.addEventListener("abort", () => {
  34. Session.abort(session.id)
  35. })
  36. const result = await Session.chat({
  37. messageID,
  38. sessionID: session.id,
  39. modelID: msg.modelID,
  40. providerID: msg.providerID,
  41. parts: [
  42. {
  43. id: Identifier.ascending("part"),
  44. messageID,
  45. sessionID: session.id,
  46. type: "text",
  47. text: params.prompt,
  48. },
  49. ],
  50. })
  51. unsub()
  52. return {
  53. title: params.description,
  54. metadata: {
  55. summary: result.parts.filter((x) => x.type === "tool"),
  56. },
  57. output: result.parts.findLast((x) => x.type === "text")!.text,
  58. }
  59. },
  60. })