tool.ts 2.9 KB

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