task.ts 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  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 { Message } from "../session/message"
  7. export const TaskTool = Tool.define({
  8. id: "task",
  9. description: DESCRIPTION,
  10. parameters: z.object({
  11. description: z
  12. .string()
  13. .describe("A short (3-5 words) description of the task"),
  14. prompt: z.string().describe("The task for the agent to perform"),
  15. }),
  16. async execute(params, ctx) {
  17. const session = await Session.create(ctx.sessionID)
  18. const msg = await Session.getMessage(ctx.sessionID, ctx.messageID)
  19. const metadata = msg.metadata.assistant!
  20. function summary(input: Message.Info) {
  21. const result = []
  22. for (const part of input.parts) {
  23. if (part.type === "tool-invocation") {
  24. result.push({
  25. toolInvocation: part.toolInvocation,
  26. metadata: input.metadata.tool[part.toolInvocation.toolCallId],
  27. })
  28. }
  29. }
  30. return result
  31. }
  32. const unsub = Bus.subscribe(Message.Event.Updated, async (evt) => {
  33. if (evt.properties.info.metadata.sessionID !== session.id) return
  34. ctx.metadata({
  35. title: params.description,
  36. summary: summary(evt.properties.info),
  37. })
  38. })
  39. ctx.abort.addEventListener("abort", () => {
  40. Session.abort(session.id)
  41. })
  42. const result = await Session.chat({
  43. sessionID: session.id,
  44. modelID: metadata.modelID,
  45. providerID: metadata.providerID,
  46. parts: [
  47. {
  48. type: "text",
  49. text: params.prompt,
  50. },
  51. ],
  52. })
  53. unsub()
  54. return {
  55. metadata: {
  56. title: params.description,
  57. summary: summary(result),
  58. },
  59. output: result.parts.findLast((x) => x.type === "text")!.text,
  60. }
  61. },
  62. })