tool.ts 3.0 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. import z from "zod"
  2. import type { MessageV2 } from "../session/message-v2"
  3. import type { Agent } from "../agent/agent"
  4. import type { PermissionNext } from "../permission/next"
  5. import { Truncate } from "./truncation"
  6. export namespace Tool {
  7. interface Metadata {
  8. [key: string]: any
  9. }
  10. export interface InitContext {
  11. agent?: Agent.Info
  12. }
  13. export type Context<M extends Metadata = Metadata> = {
  14. sessionID: string
  15. messageID: string
  16. agent: string
  17. abort: AbortSignal
  18. callID?: string
  19. extra?: { [key: string]: any }
  20. messages: MessageV2.WithParts[]
  21. metadata(input: { title?: string; metadata?: M }): void
  22. ask(input: Omit<PermissionNext.Request, "id" | "sessionID" | "tool">): Promise<void>
  23. }
  24. export interface Info<Parameters extends z.ZodType = z.ZodType, M extends Metadata = Metadata> {
  25. id: string
  26. init: (ctx?: InitContext) => Promise<{
  27. description: string
  28. parameters: Parameters
  29. execute(
  30. args: z.infer<Parameters>,
  31. ctx: Context,
  32. ): Promise<{
  33. title: string
  34. metadata: M
  35. output: string
  36. attachments?: Omit<MessageV2.FilePart, "id" | "sessionID" | "messageID">[]
  37. }>
  38. formatValidationError?(error: z.ZodError): string
  39. }>
  40. }
  41. export type InferParameters<T extends Info> = T extends Info<infer P> ? z.infer<P> : never
  42. export type InferMetadata<T extends Info> = T extends Info<any, infer M> ? M : never
  43. export function define<Parameters extends z.ZodType, Result extends Metadata>(
  44. id: string,
  45. init: Info<Parameters, Result>["init"] | Awaited<ReturnType<Info<Parameters, Result>["init"]>>,
  46. ): Info<Parameters, Result> {
  47. return {
  48. id,
  49. init: async (initCtx) => {
  50. const toolInfo = init instanceof Function ? await init(initCtx) : init
  51. const execute = toolInfo.execute
  52. toolInfo.execute = async (args, ctx) => {
  53. try {
  54. toolInfo.parameters.parse(args)
  55. } catch (error) {
  56. if (error instanceof z.ZodError && toolInfo.formatValidationError) {
  57. throw new Error(toolInfo.formatValidationError(error), { cause: error })
  58. }
  59. throw new Error(
  60. `The ${id} tool was called with invalid arguments: ${error}.\nPlease rewrite the input so it satisfies the expected schema.`,
  61. { cause: error },
  62. )
  63. }
  64. const result = await execute(args, ctx)
  65. // skip truncation for tools that handle it themselves
  66. if (result.metadata.truncated !== undefined) {
  67. return result
  68. }
  69. const truncated = await Truncate.output(result.output, {}, initCtx?.agent)
  70. return {
  71. ...result,
  72. output: truncated.content,
  73. metadata: {
  74. ...result.metadata,
  75. truncated: truncated.truncated,
  76. ...(truncated.truncated && { outputPath: truncated.outputPath }),
  77. },
  78. }
  79. }
  80. return toolInfo
  81. },
  82. }
  83. }
  84. }