@genkit-ai/googleai/lib/embedder.mjs.map

Genkit AI framework plugin for Google AI APIs, including Gemini APIs.

{"version":3,"sources":["../src/embedder.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  GoogleGenerativeAI,\n  type EmbedContentRequest,\n} from '@google/generative-ai';\nimport {\n  GenkitError,\n  z,\n  type EmbedderAction,\n  type EmbedderReference,\n  type Genkit,\n} from 'genkit';\nimport { embedderRef } from 'genkit/embedder';\nimport { getApiKeyFromEnvVar } from './common.js';\nimport type { PluginOptions } from './index.js';\n\nexport const TaskTypeSchema = z.enum([\n  'RETRIEVAL_DOCUMENT',\n  'RETRIEVAL_QUERY',\n  'SEMANTIC_SIMILARITY',\n  'CLASSIFICATION',\n  'CLUSTERING',\n]);\nexport type TaskType = z.infer<typeof TaskTypeSchema>;\n\nexport const GeminiEmbeddingConfigSchema = z.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, the model generates embeddings with 768 dimensions. Models such as\n   * `text-embedding-004`, `text-embedding-005`, and `text-multilingual-embedding-002`\n   * allow the output dimensionality to be adjusted between 1 and 768.\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).max(768).optional(),\n});\n\nexport type GeminiEmbeddingConfig = z.infer<typeof GeminiEmbeddingConfigSchema>;\n\nexport const textEmbeddingGecko001 = embedderRef({\n  name: 'googleai/embedding-001',\n  configSchema: GeminiEmbeddingConfigSchema,\n  info: {\n    dimensions: 768,\n    label: 'Google Gen AI - Text Embedding Gecko (Legacy)',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const textEmbedding004 = embedderRef({\n  name: 'googleai/text-embedding-004',\n  configSchema: GeminiEmbeddingConfigSchema,\n  info: {\n    dimensions: 768,\n    label: 'Google Gen AI - Text Embedding 001',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const SUPPORTED_MODELS = {\n  'embedding-001': textEmbeddingGecko001,\n  'text-embedding-004': textEmbedding004,\n};\n\nexport function defineGoogleAIEmbedder(\n  ai: Genkit,\n  name: string,\n  pluginOptions: PluginOptions\n): EmbedderAction<any> {\n  let apiKey: string | undefined;\n  // DO NOT throw if {apiKey: false} was supplied to options.\n  if (pluginOptions.apiKey !== false) {\n    apiKey = pluginOptions?.apiKey || getApiKeyFromEnvVar();\n    if (!apiKey)\n      throw new Error(\n        'Please pass in the API key or set either GEMINI_API_KEY or GOOGLE_API_KEY environment variable.\\n' +\n          'For more details see https://genkit.dev/docs/plugins/google-genai'\n      );\n  }\n  const embedder: EmbedderReference =\n    SUPPORTED_MODELS[name] ??\n    embedderRef({\n      name: name,\n      configSchema: GeminiEmbeddingConfigSchema,\n      info: {\n        dimensions: 768,\n        label: `Google AI - ${name}`,\n        supports: {\n          input: ['text', 'image', 'video'],\n        },\n      },\n    });\n  const apiModelName = embedder.name.startsWith('googleai/')\n    ? embedder.name.substring('googleai/'.length)\n    : embedder.name;\n  return ai.defineEmbedder(\n    {\n      name: embedder.name,\n      configSchema: GeminiEmbeddingConfigSchema,\n      info: embedder.info!,\n    },\n    async (input, options) => {\n      if (pluginOptions.apiKey === false && !options?.apiKey) {\n        throw new GenkitError({\n          status: 'INVALID_ARGUMENT',\n          message:\n            'GoogleAI plugin was initialized with {apiKey: false} but no apiKey configuration was passed at call time.',\n        });\n      }\n      const client = new GoogleGenerativeAI(\n        options?.apiKey || apiKey!\n      ).getGenerativeModel({\n        model:\n          options?.version ||\n          embedder.config?.version ||\n          embedder.version ||\n          apiModelName,\n      });\n      const embeddings = await Promise.all(\n        input.map(async (doc) => {\n          const response = await client.embedContent({\n            taskType: options?.taskType,\n            title: options?.title,\n            content: {\n              role: '',\n              parts: [{ text: doc.text }],\n            },\n            outputDimensionality: options?.outputDimensionality,\n          } as EmbedContentRequest);\n          const values = response.embedding.values;\n          return { embedding: values };\n        })\n      );\n      return { embeddings };\n    }\n  );\n}\n"],"mappings":"AAgBA;AAAA,EACE;AAAA,OAEK;AACP;AAAA,EACE;AAAA,EACA;AAAA,OAIK;AACP,SAAS,mBAAmB;AAC5B,SAAS,2BAA2B;AAG7B,MAAM,iBAAiB,EAAE,KAAK;AAAA,EACnC;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,EACA;AACF,CAAC;AAGM,MAAM,8BAA8B,EAAE,OAAO;AAAA;AAAA,EAElD,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;AAAA;AAAA,EAQ7B,sBAAsB,EAAE,OAAO,EAAE,IAAI,CAAC,EAAE,IAAI,GAAG,EAAE,SAAS;AAC5D,CAAC;AAIM,MAAM,wBAAwB,YAAY;AAAA,EAC/C,MAAM;AAAA,EACN,cAAc;AAAA,EACd,MAAM;AAAA,IACJ,YAAY;AAAA,IACZ,OAAO;AAAA,IACP,UAAU;AAAA,MACR,OAAO,CAAC,MAAM;AAAA,IAChB;AAAA,EACF;AACF,CAAC;AAEM,MAAM,mBAAmB,YAAY;AAAA,EAC1C,MAAM;AAAA,EACN,cAAc;AAAA,EACd,MAAM;AAAA,IACJ,YAAY;AAAA,IACZ,OAAO;AAAA,IACP,UAAU;AAAA,MACR,OAAO,CAAC,MAAM;AAAA,IAChB;AAAA,EACF;AACF,CAAC;AAEM,MAAM,mBAAmB;AAAA,EAC9B,iBAAiB;AAAA,EACjB,sBAAsB;AACxB;AAEO,SAAS,uBACd,IACA,MACA,eACqB;AACrB,MAAI;AAEJ,MAAI,cAAc,WAAW,OAAO;AAClC,aAAS,eAAe,UAAU,oBAAoB;AACtD,QAAI,CAAC;AACH,YAAM,IAAI;AAAA,QACR;AAAA,MAEF;AAAA,EACJ;AACA,QAAM,WACJ,iBAAiB,IAAI,KACrB,YAAY;AAAA,IACV;AAAA,IACA,cAAc;AAAA,IACd,MAAM;AAAA,MACJ,YAAY;AAAA,MACZ,OAAO,eAAe,IAAI;AAAA,MAC1B,UAAU;AAAA,QACR,OAAO,CAAC,QAAQ,SAAS,OAAO;AAAA,MAClC;AAAA,IACF;AAAA,EACF,CAAC;AACH,QAAM,eAAe,SAAS,KAAK,WAAW,WAAW,IACrD,SAAS,KAAK,UAAU,YAAY,MAAM,IAC1C,SAAS;AACb,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,SAAS;AAAA,MACf,cAAc;AAAA,MACd,MAAM,SAAS;AAAA,IACjB;AAAA,IACA,OAAO,OAAO,YAAY;AACxB,UAAI,cAAc,WAAW,SAAS,CAAC,SAAS,QAAQ;AACtD,cAAM,IAAI,YAAY;AAAA,UACpB,QAAQ;AAAA,UACR,SACE;AAAA,QACJ,CAAC;AAAA,MACH;AACA,YAAM,SAAS,IAAI;AAAA,QACjB,SAAS,UAAU;AAAA,MACrB,EAAE,mBAAmB;AAAA,QACnB,OACE,SAAS,WACT,SAAS,QAAQ,WACjB,SAAS,WACT;AAAA,MACJ,CAAC;AACD,YAAM,aAAa,MAAM,QAAQ;AAAA,QAC/B,MAAM,IAAI,OAAO,QAAQ;AACvB,gBAAM,WAAW,MAAM,OAAO,aAAa;AAAA,YACzC,UAAU,SAAS;AAAA,YACnB,OAAO,SAAS;AAAA,YAChB,SAAS;AAAA,cACP,MAAM;AAAA,cACN,OAAO,CAAC,EAAE,MAAM,IAAI,KAAK,CAAC;AAAA,YAC5B;AAAA,YACA,sBAAsB,SAAS;AAAA,UACjC,CAAwB;AACxB,gBAAM,SAAS,SAAS,UAAU;AAClC,iBAAO,EAAE,WAAW,OAAO;AAAA,QAC7B,CAAC;AAAA,MACH;AACA,aAAO,EAAE,WAAW;AAAA,IACtB;AAAA,EACF;AACF;","names":[]}