{"version":3,"sources":["../src/tool.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport {\n  action,\n  Action,\n  ActionContext,\n  ActionRunOptions,\n  assertUnstable,\n  defineAction,\n  JSONSchema7,\n  stripUndefinedProps,\n  z,\n} from '@genkit-ai/core';\nimport { HasRegistry, Registry } from '@genkit-ai/core/registry';\nimport { parseSchema, toJsonSchema } from '@genkit-ai/core/schema';\nimport { setCustomMetadataAttributes } from '@genkit-ai/core/tracing';\nimport {\n  Part,\n  ToolDefinition,\n  ToolRequestPart,\n  ToolResponsePart,\n} from './model.js';\nimport { ExecutablePrompt } from './prompt.js';\n\n/**\n * An action with a `tool` type.\n */\nexport type ToolAction<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n> = Action<I, O, z.ZodTypeAny, ToolRunOptions> & {\n  __action: {\n    metadata: {\n      type: 'tool';\n    };\n  };\n  /**\n   * respond constructs a tool response corresponding to the provided interrupt tool request\n   * using the provided reply data, validating it against the output schema of the tool if\n   * it exists.\n   *\n   * @beta\n   */\n  respond(\n    /** The interrupt tool request to which you want to respond. */\n    interrupt: ToolRequestPart,\n    /**\n     * The data with which you want to respond. Must conform to a tool's output schema or an\n     * interrupt's input schema.\n     **/\n    outputData: z.infer<O>,\n    options?: { metadata?: Record<string, any> }\n  ): ToolResponsePart;\n\n  /**\n   * restart constructs a tool request corresponding to the provided interrupt tool request\n   * that will then re-trigger the tool after e.g. a user confirms. The `resumedMetadata`\n   * supplied to this method will be passed to the tool to allow for custom handling of\n   * restart logic.\n   *\n   * @param interrupt The interrupt tool request you want to restart.\n   * @param resumedMetadata The metadata you want to provide to the tool to aide in reprocessing. Defaults to `true` if none is supplied.\n   * @param options Additional options for restarting the tool.\n   *\n   * @beta\n   */\n  restart(\n    interrupt: ToolRequestPart,\n    resumedMetadata?: any,\n    options?: {\n      /**\n       * Replace the existing input arguments to the tool with different ones, for example\n       * if the user revised an action before confirming. When input is replaced, the existing\n       * tool request will be amended in the message history.\n       **/\n      replaceInput?: z.infer<I>;\n    }\n  ): ToolRequestPart;\n};\n\nexport interface ToolRunOptions extends ActionRunOptions<z.ZodTypeAny> {\n  /**\n   * If resumed is supplied to a tool at runtime, that means that it was previously interrupted and this is a second\n   * @beta\n   **/\n  resumed?: boolean | Record<string, any>;\n  /** The metadata from the tool request that triggered this run. */\n  metadata?: Record<string, any>;\n}\n\n/**\n * Configuration for a tool.\n */\nexport interface ToolConfig<I extends z.ZodTypeAny, O extends z.ZodTypeAny> {\n  /** Unique name of the tool to use as a key in the registry. */\n  name: string;\n  /** Description of the tool. This is passed to the model to help understand what the tool is used for. */\n  description: string;\n  /** Input Zod schema. Mutually exclusive with `inputJsonSchema`. */\n  inputSchema?: I;\n  /** Input JSON schema. Mutually exclusive with `inputSchema`. */\n  inputJsonSchema?: JSONSchema7;\n  /** Output Zod schema. Mutually exclusive with `outputJsonSchema`. */\n  outputSchema?: O;\n  /** Output JSON schema. Mutually exclusive with `outputSchema`. */\n  outputJsonSchema?: JSONSchema7;\n  /** Metadata to be passed to the tool. */\n  metadata?: Record<string, any>;\n}\n\n/**\n * A reference to a tool in the form of a name, definition, or the action itself.\n */\nexport type ToolArgument<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n> = string | ToolAction<I, O> | Action<I, O> | ExecutablePrompt<any, any, any>;\n\n/**\n * Converts an action to a tool action by setting the appropriate metadata.\n */\nexport function asTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n  registry: Registry,\n  action: Action<I, O>\n): ToolAction<I, O> {\n  if (action.__action?.metadata?.type === 'tool') {\n    return action as ToolAction<I, O>;\n  }\n\n  const fn = ((input) => {\n    setCustomMetadataAttributes(registry, { subtype: 'tool' });\n    return action(input);\n  }) as ToolAction<I, O>;\n  fn.__action = {\n    ...action.__action,\n    metadata: { ...action.__action.metadata, type: 'tool' },\n  };\n  return fn;\n}\n\n/**\n * Resolves a mix of various formats of tool references to a list of tool actions by looking them up in the registry.\n */\nexport async function resolveTools<\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  registry: Registry,\n  tools?: (ToolArgument | ToolDefinition)[]\n): Promise<ToolAction[]> {\n  if (!tools || tools.length === 0) {\n    return [];\n  }\n\n  return await Promise.all(\n    tools.map(async (ref): Promise<ToolAction> => {\n      if (typeof ref === 'string') {\n        return await lookupToolByName(registry, ref);\n      } else if ((ref as Action).__action) {\n        return asTool(registry, ref as Action);\n      } else if (typeof (ref as ExecutablePrompt).asTool === 'function') {\n        return await (ref as ExecutablePrompt).asTool();\n      } else if (ref.name) {\n        return await lookupToolByName(\n          registry,\n          (ref as ToolDefinition).metadata?.originalName || ref.name\n        );\n      }\n      throw new Error('Tools must be strings, tool definitions, or actions.');\n    })\n  );\n}\n\nexport async function lookupToolByName(\n  registry: Registry,\n  name: string\n): Promise<ToolAction> {\n  let tool =\n    (await registry.lookupAction(name)) ||\n    (await registry.lookupAction(`/tool/${name}`)) ||\n    (await registry.lookupAction(`/prompt/${name}`));\n  if (!tool) {\n    throw new Error(`Tool ${name} not found`);\n  }\n  return tool as ToolAction;\n}\n\n/**\n * Converts a tool action to a definition of the tool to be passed to a model.\n */\nexport function toToolDefinition(\n  tool: Action<z.ZodTypeAny, z.ZodTypeAny>\n): ToolDefinition {\n  const originalName = tool.__action.name;\n  let name = originalName;\n  if (originalName.includes('/')) {\n    name = originalName.substring(originalName.lastIndexOf('/') + 1);\n  }\n\n  const out: ToolDefinition = {\n    name,\n    description: tool.__action.description || '',\n    outputSchema: toJsonSchema({\n      schema: tool.__action.outputSchema ?? z.void(),\n      jsonSchema: tool.__action.outputJsonSchema,\n    })!,\n    inputSchema: toJsonSchema({\n      schema: tool.__action.inputSchema ?? z.void(),\n      jsonSchema: tool.__action.inputJsonSchema,\n    })!,\n  };\n\n  if (originalName !== name) {\n    out.metadata = { originalName };\n  }\n\n  return out;\n}\n\nexport interface ToolFnOptions {\n  /**\n   * A function that can be called during tool execution that will result in the tool\n   * getting interrupted (immediately) and tool request returned to the upstream caller.\n   */\n  interrupt: (metadata?: Record<string, any>) => never;\n\n  context: ActionContext;\n}\n\nexport type ToolFn<I extends z.ZodTypeAny, O extends z.ZodTypeAny> = (\n  input: z.infer<I>,\n  ctx: ToolFnOptions & ToolRunOptions\n) => Promise<z.infer<O>>;\n\n/**\n * Defines a tool.\n *\n * A tool is an action that can be passed to a model to be called automatically if it so chooses.\n */\nexport function defineTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n  registry: Registry,\n  config: ToolConfig<I, O>,\n  fn: ToolFn<I, O>\n): ToolAction<I, O> {\n  const a = defineAction(\n    registry,\n    {\n      ...config,\n      actionType: 'tool',\n      metadata: { ...(config.metadata || {}), type: 'tool' },\n    },\n    (i, runOptions) => {\n      return fn(i, {\n        ...runOptions,\n        context: { ...runOptions.context },\n        interrupt: interruptTool(registry),\n      });\n    }\n  );\n  implementTool(a as ToolAction<I, O>, config, registry);\n  return a as ToolAction<I, O>;\n}\n\nfunction implementTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n  a: ToolAction<I, O>,\n  config: ToolConfig<I, O>,\n  registry: Registry\n) {\n  (a as ToolAction<I, O>).respond = (interrupt, responseData, options) => {\n    assertUnstable(\n      registry,\n      'beta',\n      \"The 'tool.reply' method is part of the 'interrupts' beta feature.\"\n    );\n    parseSchema(responseData, {\n      jsonSchema: config.outputJsonSchema,\n      schema: config.outputSchema,\n    });\n    return {\n      toolResponse: stripUndefinedProps({\n        name: interrupt.toolRequest.name,\n        ref: interrupt.toolRequest.ref,\n        output: responseData,\n      }),\n      metadata: {\n        interruptResponse: options?.metadata || true,\n      },\n    };\n  };\n\n  (a as ToolAction<I, O>).restart = (interrupt, resumedMetadata, options) => {\n    assertUnstable(\n      registry,\n      'beta',\n      \"The 'tool.restart' method is part of the 'interrupts' beta feature.\"\n    );\n    let replaceInput = options?.replaceInput;\n    if (replaceInput) {\n      replaceInput = parseSchema(replaceInput, {\n        schema: config.inputSchema,\n        jsonSchema: config.inputJsonSchema,\n      });\n    }\n    return {\n      toolRequest: stripUndefinedProps({\n        name: interrupt.toolRequest.name,\n        ref: interrupt.toolRequest.ref,\n        input: replaceInput || interrupt.toolRequest.input,\n      }),\n      metadata: stripUndefinedProps({\n        ...interrupt.metadata,\n        resumed: resumedMetadata || true,\n        // annotate the original input if replacing it\n        replacedInput: replaceInput ? interrupt.toolRequest.input : undefined,\n      }),\n    };\n  };\n}\n\n/** InterruptConfig defines the options for configuring an interrupt. */\nexport type InterruptConfig<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  R extends z.ZodTypeAny = z.ZodTypeAny,\n> = ToolConfig<I, R> & {\n  /** requestMetadata adds additional `interrupt` metadata to the `toolRequest` generated by the interrupt */\n  requestMetadata?:\n    | Record<string, any>\n    | ((\n        input: z.infer<I>\n      ) => Record<string, any> | Promise<Record<string, any>>);\n};\n\nexport function isToolRequest(part: Part): part is ToolRequestPart {\n  return !!part.toolRequest;\n}\n\nexport function isToolResponse(part: Part): part is ToolResponsePart {\n  return !!part.toolResponse;\n}\n\nexport function defineInterrupt<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n  registry: Registry,\n  config: InterruptConfig<I, O>\n): ToolAction<I, O> {\n  const { requestMetadata, ...toolConfig } = config;\n\n  return defineTool<I, O>(\n    registry,\n    toolConfig,\n    async (input, { interrupt }) => {\n      if (!config.requestMetadata) interrupt();\n      else if (typeof config.requestMetadata === 'object')\n        interrupt(config.requestMetadata);\n      else interrupt(await Promise.resolve(config.requestMetadata(input)));\n    }\n  );\n}\n\n/**\n * Thrown when tools execution is interrupted. It's meant to be caugh by the framework, not public API.\n */\nexport class ToolInterruptError extends Error {\n  constructor(readonly metadata?: Record<string, any>) {\n    super();\n    this.name = 'ToolInterruptError';\n  }\n}\n\n/**\n * Interrupts current tool execution causing tool request to be returned in the generation response.\n * Should only be called within a tool.\n */\nfunction interruptTool(registry: Registry) {\n  return (metadata?: Record<string, any>): never => {\n    assertUnstable(registry, 'beta', 'Tool interrupts are a beta feature.');\n    throw new ToolInterruptError(metadata);\n  };\n}\n\n/**\n * Defines a dynamic tool. Dynamic tools are just like regular tools but will not be registered in the\n * Genkit registry and can be defined dynamically at runtime.\n */\nexport function dynamicTool<I extends z.ZodTypeAny, O extends z.ZodTypeAny>(\n  ai: HasRegistry,\n  config: ToolConfig<I, O>,\n  fn?: ToolFn<I, O>\n): ToolAction<I, O> {\n  const a = action(\n    ai.registry,\n    {\n      ...config,\n      actionType: 'tool',\n      metadata: { ...(config.metadata || {}), type: 'tool', dynamic: true },\n    },\n    (i, runOptions) => {\n      const interrupt = interruptTool(ai.registry);\n      if (fn) {\n        return fn(i, {\n          ...runOptions,\n          context: { ...runOptions.context },\n          interrupt,\n        });\n      }\n      return interrupt();\n    }\n  );\n  implementTool(a as ToolAction<I, O>, config, ai.registry);\n  return a as ToolAction<I, O>;\n}\n"],"mappings":"AAgBA;AAAA,EACE;AAAA,EAIA;AAAA,EACA;AAAA,EAEA;AAAA,EACA;AAAA,OACK;AAEP,SAAS,aAAa,oBAAoB;AAC1C,SAAS,mCAAmC;AA0GrC,SAAS,OACd,UACAA,SACkB;AAClB,MAAIA,QAAO,UAAU,UAAU,SAAS,QAAQ;AAC9C,WAAOA;AAAA,EACT;AAEA,QAAM,KAAM,CAAC,UAAU;AACrB,gCAA4B,UAAU,EAAE,SAAS,OAAO,CAAC;AACzD,WAAOA,QAAO,KAAK;AAAA,EACrB;AACA,KAAG,WAAW;AAAA,IACZ,GAAGA,QAAO;AAAA,IACV,UAAU,EAAE,GAAGA,QAAO,SAAS,UAAU,MAAM,OAAO;AAAA,EACxD;AACA,SAAO;AACT;AAKA,eAAsB,aAIpB,UACA,OACuB;AACvB,MAAI,CAAC,SAAS,MAAM,WAAW,GAAG;AAChC,WAAO,CAAC;AAAA,EACV;AAEA,SAAO,MAAM,QAAQ;AAAA,IACnB,MAAM,IAAI,OAAO,QAA6B;AAC5C,UAAI,OAAO,QAAQ,UAAU;AAC3B,eAAO,MAAM,iBAAiB,UAAU,GAAG;AAAA,MAC7C,WAAY,IAAe,UAAU;AACnC,eAAO,OAAO,UAAU,GAAa;AAAA,MACvC,WAAW,OAAQ,IAAyB,WAAW,YAAY;AACjE,eAAO,MAAO,IAAyB,OAAO;AAAA,MAChD,WAAW,IAAI,MAAM;AACnB,eAAO,MAAM;AAAA,UACX;AAAA,UACC,IAAuB,UAAU,gBAAgB,IAAI;AAAA,QACxD;AAAA,MACF;AACA,YAAM,IAAI,MAAM,sDAAsD;AAAA,IACxE,CAAC;AAAA,EACH;AACF;AAEA,eAAsB,iBACpB,UACA,MACqB;AACrB,MAAI,OACD,MAAM,SAAS,aAAa,IAAI,KAChC,MAAM,SAAS,aAAa,SAAS,IAAI,EAAE,KAC3C,MAAM,SAAS,aAAa,WAAW,IAAI,EAAE;AAChD,MAAI,CAAC,MAAM;AACT,UAAM,IAAI,MAAM,QAAQ,IAAI,YAAY;AAAA,EAC1C;AACA,SAAO;AACT;AAKO,SAAS,iBACd,MACgB;AAChB,QAAM,eAAe,KAAK,SAAS;AACnC,MAAI,OAAO;AACX,MAAI,aAAa,SAAS,GAAG,GAAG;AAC9B,WAAO,aAAa,UAAU,aAAa,YAAY,GAAG,IAAI,CAAC;AAAA,EACjE;AAEA,QAAM,MAAsB;AAAA,IAC1B;AAAA,IACA,aAAa,KAAK,SAAS,eAAe;AAAA,IAC1C,cAAc,aAAa;AAAA,MACzB,QAAQ,KAAK,SAAS,gBAAgB,EAAE,KAAK;AAAA,MAC7C,YAAY,KAAK,SAAS;AAAA,IAC5B,CAAC;AAAA,IACD,aAAa,aAAa;AAAA,MACxB,QAAQ,KAAK,SAAS,eAAe,EAAE,KAAK;AAAA,MAC5C,YAAY,KAAK,SAAS;AAAA,IAC5B,CAAC;AAAA,EACH;AAEA,MAAI,iBAAiB,MAAM;AACzB,QAAI,WAAW,EAAE,aAAa;AAAA,EAChC;AAEA,SAAO;AACT;AAsBO,SAAS,WACd,UACA,QACA,IACkB;AAClB,QAAM,IAAI;AAAA,IACR;AAAA,IACA;AAAA,MACE,GAAG;AAAA,MACH,YAAY;AAAA,MACZ,UAAU,EAAE,GAAI,OAAO,YAAY,CAAC,GAAI,MAAM,OAAO;AAAA,IACvD;AAAA,IACA,CAAC,GAAG,eAAe;AACjB,aAAO,GAAG,GAAG;AAAA,QACX,GAAG;AAAA,QACH,SAAS,EAAE,GAAG,WAAW,QAAQ;AAAA,QACjC,WAAW,cAAc,QAAQ;AAAA,MACnC,CAAC;AAAA,IACH;AAAA,EACF;AACA,gBAAc,GAAuB,QAAQ,QAAQ;AACrD,SAAO;AACT;AAEA,SAAS,cACP,GACA,QACA,UACA;AACA,EAAC,EAAuB,UAAU,CAAC,WAAW,cAAc,YAAY;AACtE;AAAA,MACE;AAAA,MACA;AAAA,MACA;AAAA,IACF;AACA,gBAAY,cAAc;AAAA,MACxB,YAAY,OAAO;AAAA,MACnB,QAAQ,OAAO;AAAA,IACjB,CAAC;AACD,WAAO;AAAA,MACL,cAAc,oBAAoB;AAAA,QAChC,MAAM,UAAU,YAAY;AAAA,QAC5B,KAAK,UAAU,YAAY;AAAA,QAC3B,QAAQ;AAAA,MACV,CAAC;AAAA,MACD,UAAU;AAAA,QACR,mBAAmB,SAAS,YAAY;AAAA,MAC1C;AAAA,IACF;AAAA,EACF;AAEA,EAAC,EAAuB,UAAU,CAAC,WAAW,iBAAiB,YAAY;AACzE;AAAA,MACE;AAAA,MACA;AAAA,MACA;AAAA,IACF;AACA,QAAI,eAAe,SAAS;AAC5B,QAAI,cAAc;AAChB,qBAAe,YAAY,cAAc;AAAA,QACvC,QAAQ,OAAO;AAAA,QACf,YAAY,OAAO;AAAA,MACrB,CAAC;AAAA,IACH;AACA,WAAO;AAAA,MACL,aAAa,oBAAoB;AAAA,QAC/B,MAAM,UAAU,YAAY;AAAA,QAC5B,KAAK,UAAU,YAAY;AAAA,QAC3B,OAAO,gBAAgB,UAAU,YAAY;AAAA,MAC/C,CAAC;AAAA,MACD,UAAU,oBAAoB;AAAA,QAC5B,GAAG,UAAU;AAAA,QACb,SAAS,mBAAmB;AAAA;AAAA,QAE5B,eAAe,eAAe,UAAU,YAAY,QAAQ;AAAA,MAC9D,CAAC;AAAA,IACH;AAAA,EACF;AACF;AAeO,SAAS,cAAc,MAAqC;AACjE,SAAO,CAAC,CAAC,KAAK;AAChB;AAEO,SAAS,eAAe,MAAsC;AACnE,SAAO,CAAC,CAAC,KAAK;AAChB;AAEO,SAAS,gBACd,UACA,QACkB;AAClB,QAAM,EAAE,iBAAiB,GAAG,WAAW,IAAI;AAE3C,SAAO;AAAA,IACL;AAAA,IACA;AAAA,IACA,OAAO,OAAO,EAAE,UAAU,MAAM;AAC9B,UAAI,CAAC,OAAO,gBAAiB,WAAU;AAAA,eAC9B,OAAO,OAAO,oBAAoB;AACzC,kBAAU,OAAO,eAAe;AAAA,UAC7B,WAAU,MAAM,QAAQ,QAAQ,OAAO,gBAAgB,KAAK,CAAC,CAAC;AAAA,IACrE;AAAA,EACF;AACF;AAKO,MAAM,2BAA2B,MAAM;AAAA,EAC5C,YAAqB,UAAgC;AACnD,UAAM;AADa;AAEnB,SAAK,OAAO;AAAA,EACd;AACF;AAMA,SAAS,cAAc,UAAoB;AACzC,SAAO,CAAC,aAA0C;AAChD,mBAAe,UAAU,QAAQ,qCAAqC;AACtE,UAAM,IAAI,mBAAmB,QAAQ;AAAA,EACvC;AACF;AAMO,SAAS,YACd,IACA,QACA,IACkB;AAClB,QAAM,IAAI;AAAA,IACR,GAAG;AAAA,IACH;AAAA,MACE,GAAG;AAAA,MACH,YAAY;AAAA,MACZ,UAAU,EAAE,GAAI,OAAO,YAAY,CAAC,GAAI,MAAM,QAAQ,SAAS,KAAK;AAAA,IACtE;AAAA,IACA,CAAC,GAAG,eAAe;AACjB,YAAM,YAAY,cAAc,GAAG,QAAQ;AAC3C,UAAI,IAAI;AACN,eAAO,GAAG,GAAG;AAAA,UACX,GAAG;AAAA,UACH,SAAS,EAAE,GAAG,WAAW,QAAQ;AAAA,UACjC;AAAA,QACF,CAAC;AAAA,MACH;AACA,aAAO,UAAU;AAAA,IACnB;AAAA,EACF;AACA,gBAAc,GAAuB,QAAQ,GAAG,QAAQ;AACxD,SAAO;AACT;","names":["action"]}