{"version":3,"sources":["../../src/generate/middleware.ts"],"sourcesContent":["/**\n * Copyright 2025 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 { ActionRunOptions, GenkitError, z } from '@genkit-ai/core';\nimport type { Registry } from '@genkit-ai/core/registry';\nimport { toJsonSchema } from '@genkit-ai/core/schema';\nimport { GenkitAI } from '../genkit-ai.js';\nimport type { GenerateActionOptions } from '../model-types.js';\nimport { type MiddlewareRef } from '../model-types.js';\nimport {\n  GenerateRequest,\n  GenerateResponseChunkData,\n  GenerateResponseData,\n  ToolRequestPart,\n  ToolResponsePart,\n} from '../model.js';\nimport { type GenkitPluginV2 } from '../plugin.js';\nimport { ToolAction } from '../tool.js';\n\n/** Descriptor for a registered middleware, returned by reflection API. */\nexport const MiddlewareDescSchema = z.object({\n  /** Unique name of the middleware. */\n  name: z.string(),\n  /** Human-readable description of what the middleware does. */\n  description: z.string().optional(),\n  /** JSON Schema for the middleware's configuration. */\n  configSchema: z.record(z.any()).nullish(),\n  /** User defined metadata for the middleware. */\n  metadata: z.record(z.any()).nullish(),\n});\nexport type MiddlewareDesc = z.infer<typeof MiddlewareDescSchema>;\n\n/**\n * Defines a Genkit Generate Middleware instance, which can be configured and registered.\n * When invoked with an optional configuration, it returns a reference suitable for\n * inclusion in a `GenerateOptions.use` array.\n */\nexport interface GenerateMiddleware<\n  ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny,\n  PluginOptions = void,\n> extends MiddlewareDesc {\n  /** Configures the middleware, returning a MiddlewareRef for usage in `generate({use: [...]})`. */\n  (\n    config?: z.infer<ConfigSchema>\n  ): MiddlewareRef & { __def: GenerateMiddleware<ConfigSchema, PluginOptions> };\n  /** The unique name of this middleware. */\n  name: string;\n  /** An optional description of what the middleware does. */\n  description?: string;\n  /** An optional Zod schema for validating the middleware's configuration. */\n  configSchema?: ConfigSchema;\n  /** Metadata describing this middleware. */\n  metadata?: Record<string, any>;\n  /**\n   * Factory function that receives the validated configuration and creates\n   * a `GenerateMiddlewareDef` holding the active hooks.\n   */\n  instantiate: (\n    options: MiddlewareFnOptions<ConfigSchema, PluginOptions>\n  ) => GenerateMiddlewareDef;\n  /**\n   * Optional plugin wrapper exposing this middleware for framework-level registration.\n   */\n  plugin: (pluginOptions: PluginOptions) => GenkitPluginV2;\n  /** Generates a JSON-compatible representation of the middleware metadata. */\n  toJson: () => MiddlewareDesc;\n}\n\ninterface GenerateMiddlewarePluginInstance<\n  ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny,\n  PluginOptions = void,\n> extends GenerateMiddleware<ConfigSchema, PluginOptions> {\n  /** Options passed to the middleware plugin function. */\n  pluginOptions: PluginOptions;\n}\n\n/**\n * An instantiated implementation of a Generate Middleware.\n * Provides optional hooks to intercept the high-level `generate` action,\n * the underlying `model` execution, or individual `tool` calls, as well as\n * tools to inject into the execution.\n */\nexport interface GenerateMiddlewareDef {\n  /**\n   * Hook for intercepting the top-level generate action.\n   * Can be used to inject request parameters, modify the response, or catch errors.\n   */\n  generate?: (\n    envelope: {\n      request: GenerateActionOptions;\n      currentTurn: number;\n      messageIndex: number;\n    },\n    ctx: ActionRunOptions<GenerateResponseChunkData>,\n    next: (\n      envelope: {\n        request: GenerateActionOptions;\n        currentTurn: number;\n        messageIndex: number;\n      },\n      ctx: ActionRunOptions<GenerateResponseChunkData>\n    ) => Promise<GenerateResponseData>\n  ) => Promise<GenerateResponseData>;\n  /**\n   * Hook for intercepting the underlying model execution.\n   * Ideal for model-level caching, retry logic, or prompt/response parsing.\n   */\n  model?: (\n    req: GenerateRequest<any>,\n    ctx: ActionRunOptions<GenerateResponseChunkData>,\n    next: (\n      req: GenerateRequest<any>,\n      ctx: ActionRunOptions<GenerateResponseChunkData>\n    ) => Promise<GenerateResponseData>\n  ) => Promise<GenerateResponseData>;\n  /**\n   * Hook for intercepting individual tool calls.\n   * Enables caching tool responses, validating inputs, or overriding tool execution.\n   */\n  tool?: (\n    req: ToolRequestPart,\n    ctx: ActionRunOptions<any>,\n    next: (\n      req: ToolRequestPart,\n      ctx: ActionRunOptions<any>\n    ) => Promise<ToolResponsePart | undefined>\n  ) => Promise<ToolResponsePart | undefined>;\n  /**\n   * Tools to statically inject into the generation request whenever this middleware is active.\n   */\n  tools?: ToolAction[];\n}\n\n/**\n * Options passed to the middleware factory function.\n */\nexport interface MiddlewareFnOptions<\n  ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny,\n  PluginOptions = void,\n> {\n  config: z.infer<ConfigSchema> | undefined;\n  pluginConfig: PluginOptions;\n  ai: GenkitAI;\n}\n\nexport function generateMiddleware<\n  PluginOptions = void,\n  ConfigSchema extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  options: {\n    name: string;\n    configSchema?: ConfigSchema;\n    description?: string;\n    metadata?: Record<string, any>;\n  },\n  middlewareFn: (\n    options: MiddlewareFnOptions<ConfigSchema, PluginOptions>\n  ) => GenerateMiddlewareDef\n): GenerateMiddleware<ConfigSchema, PluginOptions> {\n  const def = function (config?: z.infer<ConfigSchema>) {\n    return {\n      name: options.name,\n      config,\n      __def: def,\n    };\n  } as GenerateMiddleware<ConfigSchema, PluginOptions>;\n  Object.defineProperty(def, 'name', { value: options.name });\n  def.configSchema = options.configSchema;\n  def.description = options.description;\n  def.metadata = options.metadata;\n  def.instantiate = middlewareFn;\n  def.plugin = (pluginConfig: PluginOptions) => ({\n    name: `middleware:${options.name}`,\n    version: 'v2',\n    middleware: () => {\n      const wrappedDef = Object.assign(\n        function (config?: z.infer<ConfigSchema>) {\n          return def(config);\n        },\n        def,\n        { pluginOptions: pluginConfig }\n      ) as GenerateMiddlewarePluginInstance<ConfigSchema, PluginOptions>;\n      Object.defineProperty(wrappedDef, 'name', { value: options.name });\n      return [wrappedDef];\n    },\n    model: (_) => {\n      throw new Error('Not supported for middleware plugins');\n    },\n  });\n\n  def.toJson = () => ({\n    name: options.name,\n    description: options.description,\n    configSchema: options.configSchema\n      ? toJsonSchema({ schema: options.configSchema })\n      : undefined,\n    metadata: options.metadata,\n  });\n\n  return def;\n}\n\nexport async function resolveMiddleware(\n  registry: Registry,\n  refs?: MiddlewareRef[]\n): Promise<GenerateMiddlewareDef[]> {\n  const result: GenerateMiddlewareDef[] = [];\n  if (!refs) return result;\n  const ai = new GenkitAI(registry);\n  for (const ref of refs) {\n    const def = await registry.lookupValue<GenerateMiddleware>(\n      'middleware',\n      ref.name\n    );\n    if (!def) {\n      throw new GenkitError({\n        status: 'NOT_FOUND',\n        message: `Middleware ${ref.name} not found in registry.`,\n      });\n    }\n    result.push(\n      def.instantiate({\n        config: ref.config,\n        ai,\n        pluginConfig: (def as GenerateMiddlewarePluginInstance).pluginOptions,\n      })\n    );\n  }\n  return result;\n}\n"],"mappings":"AAgBA,SAA2B,aAAa,SAAS;AAEjD,SAAS,oBAAoB;AAC7B,SAAS,gBAAgB;AAclB,MAAM,uBAAuB,EAAE,OAAO;AAAA;AAAA,EAE3C,MAAM,EAAE,OAAO;AAAA;AAAA,EAEf,aAAa,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA,EAEjC,cAAc,EAAE,OAAO,EAAE,IAAI,CAAC,EAAE,QAAQ;AAAA;AAAA,EAExC,UAAU,EAAE,OAAO,EAAE,IAAI,CAAC,EAAE,QAAQ;AACtC,CAAC;AAoHM,SAAS,mBAId,SAMA,cAGiD;AACjD,QAAM,MAAM,SAAU,QAAgC;AACpD,WAAO;AAAA,MACL,MAAM,QAAQ;AAAA,MACd;AAAA,MACA,OAAO;AAAA,IACT;AAAA,EACF;AACA,SAAO,eAAe,KAAK,QAAQ,EAAE,OAAO,QAAQ,KAAK,CAAC;AAC1D,MAAI,eAAe,QAAQ;AAC3B,MAAI,cAAc,QAAQ;AAC1B,MAAI,WAAW,QAAQ;AACvB,MAAI,cAAc;AAClB,MAAI,SAAS,CAAC,kBAAiC;AAAA,IAC7C,MAAM,cAAc,QAAQ,IAAI;AAAA,IAChC,SAAS;AAAA,IACT,YAAY,MAAM;AAChB,YAAM,aAAa,OAAO;AAAA,QACxB,SAAU,QAAgC;AACxC,iBAAO,IAAI,MAAM;AAAA,QACnB;AAAA,QACA;AAAA,QACA,EAAE,eAAe,aAAa;AAAA,MAChC;AACA,aAAO,eAAe,YAAY,QAAQ,EAAE,OAAO,QAAQ,KAAK,CAAC;AACjE,aAAO,CAAC,UAAU;AAAA,IACpB;AAAA,IACA,OAAO,CAAC,MAAM;AACZ,YAAM,IAAI,MAAM,sCAAsC;AAAA,IACxD;AAAA,EACF;AAEA,MAAI,SAAS,OAAO;AAAA,IAClB,MAAM,QAAQ;AAAA,IACd,aAAa,QAAQ;AAAA,IACrB,cAAc,QAAQ,eAClB,aAAa,EAAE,QAAQ,QAAQ,aAAa,CAAC,IAC7C;AAAA,IACJ,UAAU,QAAQ;AAAA,EACpB;AAEA,SAAO;AACT;AAEA,eAAsB,kBACpB,UACA,MACkC;AAClC,QAAM,SAAkC,CAAC;AACzC,MAAI,CAAC,KAAM,QAAO;AAClB,QAAM,KAAK,IAAI,SAAS,QAAQ;AAChC,aAAW,OAAO,MAAM;AACtB,UAAM,MAAM,MAAM,SAAS;AAAA,MACzB;AAAA,MACA,IAAI;AAAA,IACN;AACA,QAAI,CAAC,KAAK;AACR,YAAM,IAAI,YAAY;AAAA,QACpB,QAAQ;AAAA,QACR,SAAS,cAAc,IAAI,IAAI;AAAA,MACjC,CAAC;AAAA,IACH;AACA,WAAO;AAAA,MACL,IAAI,YAAY;AAAA,QACd,QAAQ,IAAI;AAAA,QACZ;AAAA,QACA,cAAe,IAAyC;AAAA,MAC1D,CAAC;AAAA,IACH;AAAA,EACF;AACA,SAAO;AACT;","names":[]}