UNPKG

genkitx-openai

Version:

Firebase Genkit AI framework plugin for OpenAI APIs.

1 lines 4.77 kB
{"version":3,"sources":["../src/embedder.ts"],"sourcesContent":["/**\n * Copyright 2024 The Fire Company\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\n// import { defineEmbedder, embedderRef } from '@genkit-ai/ai/embedder';\n\nimport OpenAI from 'openai';\nimport type { Genkit } from 'genkit';\nimport { embedderRef, z } from 'genkit';\n\nimport { type PluginOptions } from './index.js';\n\nexport const TextEmbeddingConfigSchema = z.object({\n dimensions: z.number().optional(),\n encodingFormat: z.union([z.literal('float'), z.literal('base64')]).optional(),\n});\n\nexport type TextEmbeddingGeckoConfig = z.infer<\n typeof TextEmbeddingConfigSchema\n>;\n\nexport const TextEmbeddingInputSchema = z.string();\n\nexport const textEmbedding3Small = embedderRef({\n name: 'openai/text-embedding-3-small',\n configSchema: TextEmbeddingConfigSchema,\n info: {\n dimensions: 1536,\n label: 'Open AI - Text Embedding 3 Small',\n supports: {\n input: ['text'],\n },\n },\n});\n\nexport const textEmbedding3Large = embedderRef({\n name: 'openai/text-embedding-3-large',\n configSchema: TextEmbeddingConfigSchema,\n info: {\n dimensions: 3072,\n label: 'Open AI - Text Embedding 3 Large',\n supports: {\n input: ['text'],\n },\n },\n});\n\nexport const textEmbeddingAda002 = embedderRef({\n name: 'openai/text-embedding-ada-002',\n configSchema: TextEmbeddingConfigSchema,\n info: {\n dimensions: 1536,\n label: 'Open AI - Text Embedding ADA 002',\n supports: {\n input: ['text'],\n },\n },\n});\n\nexport const SUPPORTED_EMBEDDING_MODELS = {\n 'text-embedding-3-small': textEmbedding3Small,\n 'text-embedding-3-large': textEmbedding3Large,\n 'text-embedding-ada-002': textEmbeddingAda002,\n};\n\nexport function openaiEmbedder(\n ai: Genkit,\n name: string,\n options?: PluginOptions\n) {\n let apiKey = options?.apiKey || process.env.OPENAI_API_KEY;\n if (!apiKey)\n throw new Error(\n 'please pass in the API key or set the OPENAI_API_KEY environment variable'\n );\n const model = SUPPORTED_EMBEDDING_MODELS[name];\n if (!model) throw new Error(`Unsupported model: ${name}`);\n\n const client = new OpenAI({ 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.embeddings.create({\n model: name,\n input: input.map((d) => d.text),\n dimensions: options?.dimensions,\n encoding_format: options?.encodingFormat,\n });\n return {\n embeddings: embeddings.data.map((d) => ({ embedding: d.embedding })),\n };\n }\n );\n}\n"],"mappings":";;;AAkBA,OAAO,YAAY;AAEnB,SAAS,aAAa,SAAS;AAIxB,MAAM,4BAA4B,EAAE,OAAO;AAAA,EAChD,YAAY,EAAE,OAAO,EAAE,SAAS;AAAA,EAChC,gBAAgB,EAAE,MAAM,CAAC,EAAE,QAAQ,OAAO,GAAG,EAAE,QAAQ,QAAQ,CAAC,CAAC,EAAE,SAAS;AAC9E,CAAC;AAMM,MAAM,2BAA2B,EAAE,OAAO;AAE1C,MAAM,sBAAsB,YAAY;AAAA,EAC7C,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,sBAAsB,YAAY;AAAA,EAC7C,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,sBAAsB,YAAY;AAAA,EAC7C,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,0BAA0B;AAAA,EAC1B,0BAA0B;AAAA,EAC1B,0BAA0B;AAC5B;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;AAExD,QAAM,SAAS,IAAI,OAAO,EAAE,OAAO,CAAC;AACpC,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,WAAW,OAAO;AAAA,QAChD,OAAO;AAAA,QACP,OAAO,MAAM,IAAI,CAAC,MAAM,EAAE,IAAI;AAAA,QAC9B,YAAYA,YAAA,gBAAAA,SAAS;AAAA,QACrB,iBAAiBA,YAAA,gBAAAA,SAAS;AAAA,MAC5B,CAAC;AACD,aAAO;AAAA,QACL,YAAY,WAAW,KAAK,IAAI,CAAC,OAAO,EAAE,WAAW,EAAE,UAAU,EAAE;AAAA,MACrE;AAAA,IACF;AAAA,EACF;AACF;","names":["options"]}