UNPKG

@langchain/community

Version:
1 lines 6.38 kB
{"version":3,"file":"deepinfra.cjs","names":["Embeddings"],"sources":["../../src/embeddings/deepinfra.ts"],"sourcesContent":["import { getEnvironmentVariable } from \"@langchain/core/utils/env\";\nimport { Embeddings, EmbeddingsParams } from \"@langchain/core/embeddings\";\nimport { chunkArray } from \"@langchain/core/utils/chunk_array\";\n\n/**\n * The default model name to use for generating embeddings.\n */\nconst DEFAULT_MODEL_NAME = \"sentence-transformers/clip-ViT-B-32\";\n\n/**\n * The default batch size to use for generating embeddings.\n * This is limited by the DeepInfra API to a maximum of 1024.\n */\nconst DEFAULT_BATCH_SIZE = 1024;\n\n/**\n * Environment variable name for the DeepInfra API token.\n */\nconst API_TOKEN_ENV_VAR = \"DEEPINFRA_API_TOKEN\";\n\nexport interface DeepInfraEmbeddingsRequest {\n inputs: string[];\n normalize?: boolean;\n image?: string;\n webhook?: string;\n}\n\n/**\n * Input parameters for the DeepInfra embeddings\n */\nexport interface DeepInfraEmbeddingsParams extends EmbeddingsParams {\n /**\n * The API token to use for authentication.\n * If not provided, it will be read from the `DEEPINFRA_API_TOKEN` environment variable.\n */\n apiToken?: string;\n\n /**\n * The model ID to use for generating completions.\n * Default: `sentence-transformers/clip-ViT-B-32`\n */\n modelName?: string;\n\n /**\n * The maximum number of texts to embed in a single request. This is\n * limited by the DeepInfra API to a maximum of 1024.\n */\n batchSize?: number;\n}\n\n/**\n * Response from the DeepInfra embeddings API.\n */\nexport interface DeepInfraEmbeddingsResponse {\n /**\n * The embeddings generated for the input texts.\n */\n embeddings: number[][];\n /**\n * The number of tokens in the input texts.\n */\n input_tokens: number;\n /**\n * The status of the inference.\n */\n request_id?: string;\n}\n\n/**\n * A class for generating embeddings using the DeepInfra API.\n * @example\n * ```typescript\n * // Embed a query using the DeepInfraEmbeddings class\n * const model = new DeepInfraEmbeddings();\n * const res = await model.embedQuery(\n * \"What would be a good company name for a company that makes colorful socks?\",\n * );\n * console.log({ res });\n * ```\n */\nexport class DeepInfraEmbeddings\n extends Embeddings\n implements DeepInfraEmbeddingsParams\n{\n apiToken: string;\n\n batchSize: number;\n\n modelName: string;\n\n /**\n * Constructor for the DeepInfraEmbeddings class.\n * @param fields - An optional object with properties to configure the instance.\n */\n constructor(\n fields?: Partial<DeepInfraEmbeddingsParams> & {\n verbose?: boolean;\n }\n ) {\n const fieldsWithDefaults = {\n modelName: DEFAULT_MODEL_NAME,\n batchSize: DEFAULT_BATCH_SIZE,\n ...fields,\n };\n\n super(fieldsWithDefaults);\n\n const apiKey =\n fieldsWithDefaults?.apiToken || getEnvironmentVariable(API_TOKEN_ENV_VAR);\n\n if (!apiKey) {\n throw new Error(\"DeepInfra API token not found\");\n }\n\n this.modelName = fieldsWithDefaults?.modelName ?? this.modelName;\n this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize;\n this.apiToken = apiKey;\n }\n\n /**\n * Generates embeddings for an array of texts.\n * @param inputs - An array of strings to generate embeddings for.\n * @returns A Promise that resolves to an array of embeddings.\n */\n async embedDocuments(inputs: string[]): Promise<number[][]> {\n const batches = chunkArray(inputs, this.batchSize);\n\n const batchRequests = batches.map((batch: string[]) =>\n this.embeddingWithRetry({\n inputs: batch,\n })\n );\n\n const batchResponses = await Promise.all(batchRequests);\n\n const out: number[][] = [];\n\n for (let i = 0; i < batchResponses.length; i += 1) {\n const batch = batches[i];\n const { embeddings } = batchResponses[i];\n for (let j = 0; j < batch.length; j += 1) {\n out.push(embeddings[j]);\n }\n }\n\n return out;\n }\n\n /**\n * Generates an embedding for a single text.\n * @param text - A string to generate an embedding for.\n * @returns A Promise that resolves to an array of numbers representing the embedding.\n */\n async embedQuery(text: string): Promise<number[]> {\n const { embeddings } = await this.embeddingWithRetry({\n inputs: [text],\n });\n return embeddings[0];\n }\n\n /**\n * Generates embeddings with retry capabilities.\n * @param request - An object containing the request parameters for generating embeddings.\n * @returns A Promise that resolves to the API response.\n */\n private async embeddingWithRetry(\n request: DeepInfraEmbeddingsRequest\n ): Promise<DeepInfraEmbeddingsResponse> {\n const response = await this.caller.call(() =>\n fetch(`https://api.deepinfra.com/v1/inference/${this.modelName}`, {\n method: \"POST\",\n headers: {\n Authorization: `Bearer ${this.apiToken}`,\n \"Content-Type\": \"application/json\",\n },\n body: JSON.stringify(request),\n }).then((res) => res.json())\n );\n return response as DeepInfraEmbeddingsResponse;\n }\n}\n"],"mappings":";;;;;;;;;;AAOA,MAAM,qBAAqB;;;;;AAM3B,MAAM,qBAAqB;;;;AAK3B,MAAM,oBAAoB;;;;;;;;;;;;;AA8D1B,IAAa,sBAAb,cACUA,2BAAAA,WAEV;CACE;CAEA;CAEA;;;;;CAMA,YACE,QAGA;EACA,MAAM,qBAAqB;GACzB,WAAW;GACX,WAAW;GACX,GAAG;GACJ;AAED,QAAM,mBAAmB;EAEzB,MAAM,SACJ,oBAAoB,aAAA,GAAA,0BAAA,wBAAmC,kBAAkB;AAE3E,MAAI,CAAC,OACH,OAAM,IAAI,MAAM,gCAAgC;AAGlD,OAAK,YAAY,oBAAoB,aAAa,KAAK;AACvD,OAAK,YAAY,oBAAoB,aAAa,KAAK;AACvD,OAAK,WAAW;;;;;;;CAQlB,MAAM,eAAe,QAAuC;EAC1D,MAAM,WAAA,GAAA,kCAAA,YAAqB,QAAQ,KAAK,UAAU;EAElD,MAAM,gBAAgB,QAAQ,KAAK,UACjC,KAAK,mBAAmB,EACtB,QAAQ,OACT,CAAC,CACH;EAED,MAAM,iBAAiB,MAAM,QAAQ,IAAI,cAAc;EAEvD,MAAM,MAAkB,EAAE;AAE1B,OAAK,IAAI,IAAI,GAAG,IAAI,eAAe,QAAQ,KAAK,GAAG;GACjD,MAAM,QAAQ,QAAQ;GACtB,MAAM,EAAE,eAAe,eAAe;AACtC,QAAK,IAAI,IAAI,GAAG,IAAI,MAAM,QAAQ,KAAK,EACrC,KAAI,KAAK,WAAW,GAAG;;AAI3B,SAAO;;;;;;;CAQT,MAAM,WAAW,MAAiC;EAChD,MAAM,EAAE,eAAe,MAAM,KAAK,mBAAmB,EACnD,QAAQ,CAAC,KAAK,EACf,CAAC;AACF,SAAO,WAAW;;;;;;;CAQpB,MAAc,mBACZ,SACsC;AAWtC,SAViB,MAAM,KAAK,OAAO,WACjC,MAAM,0CAA0C,KAAK,aAAa;GAChE,QAAQ;GACR,SAAS;IACP,eAAe,UAAU,KAAK;IAC9B,gBAAgB;IACjB;GACD,MAAM,KAAK,UAAU,QAAQ;GAC9B,CAAC,CAAC,MAAM,QAAQ,IAAI,MAAM,CAAC,CAC7B"}