{"version":3,"sources":["../../src/googleai/embedder.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 {\n  ActionMetadata,\n  EmbedderAction,\n  embedderActionMetadata,\n  EmbedderInfo,\n  EmbedderReference,\n  z,\n} from 'genkit';\nimport { embedderRef } from 'genkit/embedder';\nimport { embedder as pluginEmbedder } from 'genkit/plugin';\nimport { toGeminiMessage } from '../common/converters.mjs';\nimport { isKnownKey } from '../common/utils.mjs';\nimport { embedContent } from './client.mjs';\nimport {\n  ClientOptions,\n  EmbedContentRequest,\n  GoogleAIPluginOptions,\n  Model,\n  TaskTypeSchema,\n} from './types.mjs';\nimport {\n  calculateApiKey,\n  checkApiKey,\n  checkModelName,\n  extractVersion,\n} from './utils.mjs';\n\nexport const EmbeddingConfigSchema = z\n  .object({\n    /** Override the API key provided at plugin initialization. */\n    apiKey: z.string().optional(),\n    /**\n     * The `task_type` parameter is defined as the intended downstream application to help the model\n     * produce better quality embeddings.\n     **/\n    taskType: TaskTypeSchema.optional(),\n    title: z.string().optional(),\n    version: z.string().optional(),\n    /**\n     * The `outputDimensionality` parameter allows you to specify the dimensionality of the embedding output.\n     * By default, `gemini-embedding-2` and `gemini-embedding-001` generate embeddings with 3072 dimensions.\n     * By selecting a smaller output dimensionality, users can save memory and storage space, leading to more efficient computations.\n     **/\n    outputDimensionality: z.number().min(1).optional(),\n  })\n  .passthrough();\nexport type EmbeddingConfigSchemaType = typeof EmbeddingConfigSchema;\nexport type EmbeddingConfig = z.infer<EmbeddingConfigSchemaType>;\n\n// This contains all the embedder config schema types\ntype ConfigSchemaType = EmbeddingConfigSchemaType;\n\nfunction commonRef(\n  name: string,\n  info?: EmbedderInfo,\n  configSchema: ConfigSchemaType = EmbeddingConfigSchema\n): EmbedderReference<ConfigSchemaType> {\n  return embedderRef({\n    name: `googleai/${name}`,\n    configSchema,\n    info: info ?? {\n      dimensions: 3072,\n      supports: {\n        input: ['text'],\n      },\n    },\n  });\n}\n\nconst GENERIC_MODEL = commonRef('embedder');\n\n// Models we no longer suggest, but that still work: they still resolve and\n// are still registered as actions, but aren't in the KnownModels type, so\n// they don't show up in autocomplete. Use this for models that are deprecated\n// or < 1 month from shutdown (note the shutdown date).\nconst HIDDEN_MODELS = {} as const;\n\nconst KNOWN_MODELS = {\n  'gemini-embedding-001': commonRef('gemini-embedding-001'),\n  'gemini-embedding-2': commonRef('gemini-embedding-2', {\n    supports: {\n      input: ['text', 'image', 'video'],\n    },\n  }),\n} as const;\n\nconst ALL_MODELS = {\n  ...HIDDEN_MODELS,\n  ...KNOWN_MODELS,\n};\nexport type KnownModels = keyof typeof KNOWN_MODELS; // For autocomplete\n\nexport type EmbedderModelName = `gemini-embedding-${string}`;\nexport function isEmbedderName(value: string): value is EmbedderModelName {\n  return value.startsWith('gemini-embedding-');\n}\n\nexport function model(\n  version: string,\n  config: EmbeddingConfig = {}\n): EmbedderReference<ConfigSchemaType> {\n  const name = checkModelName(version);\n\n  if (isKnownKey(name, ALL_MODELS)) {\n    const known = ALL_MODELS[name];\n    return embedderRef({\n      name: known.name,\n      info: known.info,\n      configSchema: known.configSchema,\n      config,\n    });\n  }\n\n  return embedderRef({\n    name: `googleai/${name}`,\n    info: { ...GENERIC_MODEL.info },\n    configSchema: GENERIC_MODEL.configSchema,\n    config,\n  });\n}\n\nexport function listActions(models: Model[]): ActionMetadata[] {\n  return (\n    models\n      .filter((m) => m.supportedGenerationMethods.includes('embedContent'))\n      // Filter out deprecated\n      .filter((m) => !m.description || !m.description.includes('deprecated'))\n      .map((m) => {\n        const ref = model(m.name);\n        return embedderActionMetadata({\n          name: ref.name,\n          info: ref.info,\n          configSchema: ref.configSchema,\n        });\n      })\n  );\n}\n\nexport function listKnownModels(options?: GoogleAIPluginOptions) {\n  return Object.keys(ALL_MODELS).map((name) => defineEmbedder(name, options));\n}\n\nexport function defineEmbedder(\n  name: string,\n  pluginOptions?: GoogleAIPluginOptions\n): EmbedderAction {\n  checkApiKey(pluginOptions?.apiKey);\n  const ref = model(name);\n  const clientOptions: ClientOptions = {\n    apiVersion: pluginOptions?.apiVersion,\n    baseUrl: pluginOptions?.baseUrl,\n    customHeaders: pluginOptions?.customHeaders,\n    experimental_debugTraces: pluginOptions?.experimental_debugTraces,\n  };\n\n  return pluginEmbedder(\n    {\n      name: ref.name,\n      configSchema: ref.configSchema,\n      info: ref.info,\n    },\n    async (request, _) => {\n      const embedApiKey = calculateApiKey(\n        pluginOptions?.apiKey,\n        request.options?.apiKey\n      );\n      const embedVersion = request.options?.version || extractVersion(ref);\n\n      const embeddings = await Promise.all(\n        request.input.map(async (doc) => {\n          const response = await embedContent(\n            embedApiKey,\n            embedVersion,\n            {\n              taskType: request.options?.taskType,\n              title: request.options?.title,\n              content: toGeminiMessage({ role: 'user', content: doc.content }),\n              outputDimensionality: request.options?.outputDimensionality,\n            } as EmbedContentRequest,\n            clientOptions\n          );\n          const values = response.embedding.values;\n          return { embedding: values };\n        })\n      );\n      return { embeddings };\n    }\n  );\n}\n\nexport const TEST_ONLY = { KNOWN_MODELS };\n"],"mappings":"AAgBA;AAAA,EAGE;AAAA,EAGA;AAAA,OACK;AACP,SAAS,mBAAmB;AAC5B,SAAS,YAAY,sBAAsB;AAC3C,SAAS,uBAAuB;AAChC,SAAS,kBAAkB;AAC3B,SAAS,oBAAoB;AAC7B;AAAA,EAKE;AAAA,OACK;AACP;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,OACK;AAEA,MAAM,wBAAwB,EAClC,OAAO;AAAA;AAAA,EAEN,QAAQ,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA;AAAA;AAAA;AAAA,EAK5B,UAAU,eAAe,SAAS;AAAA,EAClC,OAAO,EAAE,OAAO,EAAE,SAAS;AAAA,EAC3B,SAAS,EAAE,OAAO,EAAE,SAAS;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,EAM7B,sBAAsB,EAAE,OAAO,EAAE,IAAI,CAAC,EAAE,SAAS;AACnD,CAAC,EACA,YAAY;AAOf,SAAS,UACP,MACA,MACA,eAAiC,uBACI;AACrC,SAAO,YAAY;AAAA,IACjB,MAAM,YAAY,IAAI;AAAA,IACtB;AAAA,IACA,MAAM,QAAQ;AAAA,MACZ,YAAY;AAAA,MACZ,UAAU;AAAA,QACR,OAAO,CAAC,MAAM;AAAA,MAChB;AAAA,IACF;AAAA,EACF,CAAC;AACH;AAEA,MAAM,gBAAgB,UAAU,UAAU;AAM1C,MAAM,gBAAgB,CAAC;AAEvB,MAAM,eAAe;AAAA,EACnB,wBAAwB,UAAU,sBAAsB;AAAA,EACxD,sBAAsB,UAAU,sBAAsB;AAAA,IACpD,UAAU;AAAA,MACR,OAAO,CAAC,QAAQ,SAAS,OAAO;AAAA,IAClC;AAAA,EACF,CAAC;AACH;AAEA,MAAM,aAAa;AAAA,EACjB,GAAG;AAAA,EACH,GAAG;AACL;AAIO,SAAS,eAAe,OAA2C;AACxE,SAAO,MAAM,WAAW,mBAAmB;AAC7C;AAEO,SAAS,MACd,SACA,SAA0B,CAAC,GACU;AACrC,QAAM,OAAO,eAAe,OAAO;AAEnC,MAAI,WAAW,MAAM,UAAU,GAAG;AAChC,UAAM,QAAQ,WAAW,IAAI;AAC7B,WAAO,YAAY;AAAA,MACjB,MAAM,MAAM;AAAA,MACZ,MAAM,MAAM;AAAA,MACZ,cAAc,MAAM;AAAA,MACpB;AAAA,IACF,CAAC;AAAA,EACH;AAEA,SAAO,YAAY;AAAA,IACjB,MAAM,YAAY,IAAI;AAAA,IACtB,MAAM,EAAE,GAAG,cAAc,KAAK;AAAA,IAC9B,cAAc,cAAc;AAAA,IAC5B;AAAA,EACF,CAAC;AACH;AAEO,SAAS,YAAY,QAAmC;AAC7D,SACE,OACG,OAAO,CAAC,MAAM,EAAE,2BAA2B,SAAS,cAAc,CAAC,EAEnE,OAAO,CAAC,MAAM,CAAC,EAAE,eAAe,CAAC,EAAE,YAAY,SAAS,YAAY,CAAC,EACrE,IAAI,CAAC,MAAM;AACV,UAAM,MAAM,MAAM,EAAE,IAAI;AACxB,WAAO,uBAAuB;AAAA,MAC5B,MAAM,IAAI;AAAA,MACV,MAAM,IAAI;AAAA,MACV,cAAc,IAAI;AAAA,IACpB,CAAC;AAAA,EACH,CAAC;AAEP;AAEO,SAAS,gBAAgB,SAAiC;AAC/D,SAAO,OAAO,KAAK,UAAU,EAAE,IAAI,CAAC,SAAS,eAAe,MAAM,OAAO,CAAC;AAC5E;AAEO,SAAS,eACd,MACA,eACgB;AAChB,cAAY,eAAe,MAAM;AACjC,QAAM,MAAM,MAAM,IAAI;AACtB,QAAM,gBAA+B;AAAA,IACnC,YAAY,eAAe;AAAA,IAC3B,SAAS,eAAe;AAAA,IACxB,eAAe,eAAe;AAAA,IAC9B,0BAA0B,eAAe;AAAA,EAC3C;AAEA,SAAO;AAAA,IACL;AAAA,MACE,MAAM,IAAI;AAAA,MACV,cAAc,IAAI;AAAA,MAClB,MAAM,IAAI;AAAA,IACZ;AAAA,IACA,OAAO,SAAS,MAAM;AACpB,YAAM,cAAc;AAAA,QAClB,eAAe;AAAA,QACf,QAAQ,SAAS;AAAA,MACnB;AACA,YAAM,eAAe,QAAQ,SAAS,WAAW,eAAe,GAAG;AAEnE,YAAM,aAAa,MAAM,QAAQ;AAAA,QAC/B,QAAQ,MAAM,IAAI,OAAO,QAAQ;AAC/B,gBAAM,WAAW,MAAM;AAAA,YACrB;AAAA,YACA;AAAA,YACA;AAAA,cACE,UAAU,QAAQ,SAAS;AAAA,cAC3B,OAAO,QAAQ,SAAS;AAAA,cACxB,SAAS,gBAAgB,EAAE,MAAM,QAAQ,SAAS,IAAI,QAAQ,CAAC;AAAA,cAC/D,sBAAsB,QAAQ,SAAS;AAAA,YACzC;AAAA,YACA;AAAA,UACF;AACA,gBAAM,SAAS,SAAS,UAAU;AAClC,iBAAO,EAAE,WAAW,OAAO;AAAA,QAC7B,CAAC;AAAA,MACH;AACA,aAAO,EAAE,WAAW;AAAA,IACtB;AAAA,EACF;AACF;AAEO,MAAM,YAAY,EAAE,aAAa;","names":[]}