{"version":3,"sources":["../src/embedder.ts"],"sourcesContent":["/**\n * Copyright 2024 Bloom Labs Inc\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 type { Genkit } from 'genkit';\nimport { z } from 'genkit';\nimport { embedderRef } from 'genkit/embedder';\nimport { CohereClient } from 'cohere-ai';\n\nimport type { PluginOptions } from '.';\n\nexport const TextEmbeddingConfigSchema = z.object({\n  // Its difficult with the schema to make this an array therefore its only a single string for now\n  // embeddingTypes: z.array(\n  //   z.enum(['float', 'int8', 'uint8', 'binary', 'ubinary'])\n  // ).optional(),\n  embeddingTypes: z\n    .union([\n      z.literal('float'),\n      z.literal('int8'),\n      z.literal('uint8'),\n      z.literal('binary'),\n      z.literal('ubinary'),\n    ])\n    .optional(),\n  inputType: z\n    .union([\n      z.literal('search_document'),\n      z.literal('search_query'),\n      z.literal('classification'),\n      z.literal('clustering'),\n    ])\n    .optional(),\n});\n\nexport type TextEmbeddingGeckoConfig = z.infer<\n  typeof TextEmbeddingConfigSchema\n>;\n\nexport const TextEmbeddingInputSchema = z.string();\n\nexport const embedMultilingual4 = embedderRef({\n  name: 'cohere/embed-v4.0',\n  configSchema: TextEmbeddingConfigSchema,\n  info: {\n    dimensions: 1024,\n    label: 'Cohere - Embed Multilingual v4.0',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const embedEnglish3 = embedderRef({\n  name: 'cohere/embed-english-v3.0',\n  configSchema: TextEmbeddingConfigSchema,\n  info: {\n    dimensions: 1024,\n    label: 'Cohere - Embed English v3.0',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const embedMultilingual3 = embedderRef({\n  name: 'cohere/embed-multilingual-v3.0',\n  configSchema: TextEmbeddingConfigSchema,\n  info: {\n    dimensions: 1024,\n    label: 'Cohere - Embed Multilingual v3.0',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const embedEnglishLight3 = embedderRef({\n  name: 'cohere/embed-english-light-v3.0',\n  configSchema: TextEmbeddingConfigSchema,\n  info: {\n    dimensions: 384,\n    label: 'Cohere - Embed English Light v3.0',\n    supports: {\n      input: ['text', 'text'],\n    },\n  },\n});\n\nexport const embedMultilingualLight3 = embedderRef({\n  name: 'cohere/embed-multilingual-light-v3.0',\n  configSchema: TextEmbeddingConfigSchema,\n  info: {\n    dimensions: 384,\n    label: 'Cohere - Embed Multilingual Light v3.0',\n    supports: {\n      input: ['text'],\n    },\n  },\n});\n\nexport const SUPPORTED_EMBEDDING_MODELS = {\n  'embed-v4.0': embedMultilingual4,\n  'embed-english-v3.0': embedEnglish3,\n  'embed-multilingual-v3.0': embedMultilingual3,\n  'embed-english-light-v3.0': embedEnglishLight3,\n  'embed-multilingual-light-v3.0': embedMultilingualLight3,\n};\n\nexport function cohereEmbedder(\n  ai: Genkit,\n  name: string,\n  options?: PluginOptions\n) {\n  let apiKey = options?.apiKey || process.env.COHERE_API_KEY;\n  if (!apiKey)\n    throw new Error(\n      'please pass in the API key or set the COHERE_API_KEY environment variable'\n    );\n  const model = SUPPORTED_EMBEDDING_MODELS[name];\n  if (!model) throw new Error(`Unsupported model: ${name}`);\n  const client = new CohereClient({ token: apiKey });\n  return ai.defineEmbedder(\n    {\n      info: model.info!,\n      configSchema: TextEmbeddingConfigSchema,\n      name: model.name,\n    },\n    async (input, options) => {\n      const embeddings = await client.embed({\n        model: name,\n        texts: input.map((d) => {\n          return d.text;\n        }),\n        inputType: options?.inputType ? options.inputType : 'search_document',\n        embeddingTypes: options?.embeddingTypes\n          ? [options.embeddingTypes]\n          : undefined,\n      });\n      try {\n        return {\n          embeddings: (embeddings.embeddings as number[][]).map((e) => ({\n            embedding: e,\n          })),\n        };\n      } catch (e) {\n        for (const property in embeddings.embeddings) {\n          if (\n            Object.prototype.hasOwnProperty.call(\n              embeddings.embeddings,\n              property\n            )\n          ) {\n            return {\n              embeddings: (embeddings.embeddings[property] as number[][]).map(\n                (e) => ({ embedding: e })\n              ),\n            };\n          }\n        }\n      }\n      return {\n        embeddings: (embeddings.embeddings as number[][]).map((e) => ({\n          embedding: e,\n        })),\n      };\n    }\n  );\n}\n"],"mappings":";;;AAiBA,SAAS,SAAS;AAClB,SAAS,mBAAmB;AAC5B,SAAS,oBAAoB;AAItB,MAAM,4BAA4B,EAAE,OAAO;AAAA;AAAA;AAAA;AAAA;AAAA,EAKhD,gBAAgB,EACb,MAAM;AAAA,IACL,EAAE,QAAQ,OAAO;AAAA,IACjB,EAAE,QAAQ,MAAM;AAAA,IAChB,EAAE,QAAQ,OAAO;AAAA,IACjB,EAAE,QAAQ,QAAQ;AAAA,IAClB,EAAE,QAAQ,SAAS;AAAA,EACrB,CAAC,EACA,SAAS;AAAA,EACZ,WAAW,EACR,MAAM;AAAA,IACL,EAAE,QAAQ,iBAAiB;AAAA,IAC3B,EAAE,QAAQ,cAAc;AAAA,IACxB,EAAE,QAAQ,gBAAgB;AAAA,IAC1B,EAAE,QAAQ,YAAY;AAAA,EACxB,CAAC,EACA,SAAS;AACd,CAAC;AAMM,MAAM,2BAA2B,EAAE,OAAO;AAE1C,MAAM,qBAAqB,YAAY;AAAA,EAC5C,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,gBAAgB,YAAY;AAAA,EACvC,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,qBAAqB,YAAY;AAAA,EAC5C,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,qBAAqB,YAAY;AAAA,EAC5C,MAAM;AAAA,EACN,cAAc;AAAA,EACd,MAAM;AAAA,IACJ,YAAY;AAAA,IACZ,OAAO;AAAA,IACP,UAAU;AAAA,MACR,OAAO,CAAC,QAAQ,MAAM;AAAA,IACxB;AAAA,EACF;AACF,CAAC;AAEM,MAAM,0BAA0B,YAAY;AAAA,EACjD,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,6BAA6B;AAAA,EACxC,cAAc;AAAA,EACd,sBAAsB;AAAA,EACtB,2BAA2B;AAAA,EAC3B,4BAA4B;AAAA,EAC5B,iCAAiC;AACnC;AAEO,SAAS,eACd,IACA,MACA,SACA;AACA,MAAI,UAAS,mCAAS,WAAU,QAAQ,IAAI;AAC5C,MAAI,CAAC;AACH,UAAM,IAAI;AAAA,MACR;AAAA,IACF;AACF,QAAM,QAAQ,2BAA2B,IAAI;AAC7C,MAAI,CAAC,MAAO,OAAM,IAAI,MAAM,sBAAsB,IAAI,EAAE;AACxD,QAAM,SAAS,IAAI,aAAa,EAAE,OAAO,OAAO,CAAC;AACjD,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,MAAM;AAAA,MACZ,cAAc;AAAA,MACd,MAAM,MAAM;AAAA,IACd;AAAA,IACA,CAAO,OAAOA,aAAY;AACxB,YAAM,aAAa,MAAM,OAAO,MAAM;AAAA,QACpC,OAAO;AAAA,QACP,OAAO,MAAM,IAAI,CAAC,MAAM;AACtB,iBAAO,EAAE;AAAA,QACX,CAAC;AAAA,QACD,YAAWA,YAAA,gBAAAA,SAAS,aAAYA,SAAQ,YAAY;AAAA,QACpD,iBAAgBA,YAAA,gBAAAA,SAAS,kBACrB,CAACA,SAAQ,cAAc,IACvB;AAAA,MACN,CAAC;AACD,UAAI;AACF,eAAO;AAAA,UACL,YAAa,WAAW,WAA0B,IAAI,CAAC,OAAO;AAAA,YAC5D,WAAW;AAAA,UACb,EAAE;AAAA,QACJ;AAAA,MACF,SAAS,GAAG;AACV,mBAAW,YAAY,WAAW,YAAY;AAC5C,cACE,OAAO,UAAU,eAAe;AAAA,YAC9B,WAAW;AAAA,YACX;AAAA,UACF,GACA;AACA,mBAAO;AAAA,cACL,YAAa,WAAW,WAAW,QAAQ,EAAiB;AAAA,gBAC1D,CAACC,QAAO,EAAE,WAAWA,GAAE;AAAA,cACzB;AAAA,YACF;AAAA,UACF;AAAA,QACF;AAAA,MACF;AACA,aAAO;AAAA,QACL,YAAa,WAAW,WAA0B,IAAI,CAAC,OAAO;AAAA,UAC5D,WAAW;AAAA,QACb,EAAE;AAAA,MACJ;AAAA,IACF;AAAA,EACF;AACF;","names":["options","e"]}