tool.ts 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import z from "zod"
  2. import type { MessageV2 } from "../session/message-v2"
  3. export namespace Tool {
  4. interface Metadata {
  5. [key: string]: any
  6. }
  7. export type Context<M extends Metadata = Metadata> = {
  8. sessionID: string
  9. messageID: string
  10. agent: string
  11. abort: AbortSignal
  12. callID?: string
  13. extra?: { [key: string]: any }
  14. metadata(input: { title?: string; metadata?: M }): void
  15. }
  16. export interface Info<Parameters extends z.ZodType = z.ZodType, M extends Metadata = Metadata> {
  17. id: string
  18. init: () => Promise<{
  19. description: string
  20. parameters: Parameters
  21. execute(
  22. args: z.infer<Parameters>,
  23. ctx: Context,
  24. ): Promise<{
  25. title: string
  26. metadata: M
  27. output: string
  28. attachments?: MessageV2.FilePart[]
  29. }>
  30. formatValidationError?(error: z.ZodError): string
  31. }>
  32. }
  33. export type InferParameters<T extends Info> = T extends Info<infer P> ? z.infer<P> : never
  34. export type InferMetadata<T extends Info> = T extends Info<any, infer M> ? M : never
  35. export function define<Parameters extends z.ZodType, Result extends Metadata>(
  36. id: string,
  37. init: Info<Parameters, Result>["init"] | Awaited<ReturnType<Info<Parameters, Result>["init"]>>,
  38. ): Info<Parameters, Result> {
  39. return {
  40. id,
  41. init: async () => {
  42. const toolInfo = init instanceof Function ? await init() : init
  43. const execute = toolInfo.execute
  44. toolInfo.execute = (args, ctx) => {
  45. try {
  46. toolInfo.parameters.parse(args)
  47. } catch (error) {
  48. if (error instanceof z.ZodError && toolInfo.formatValidationError) {
  49. throw new Error(toolInfo.formatValidationError(error))
  50. }
  51. throw error
  52. }
  53. return execute(args, ctx)
  54. }
  55. return toolInfo
  56. },
  57. }
  58. }
  59. }