UNPKG

@langchain/community

Version:
1 lines 5 kB
{"version":3,"file":"zhipuai.cjs","names":["Embeddings","encodeApiKey"],"sources":["../../src/embeddings/zhipuai.ts"],"sourcesContent":["import { getEnvironmentVariable } from \"@langchain/core/utils/env\";\n\nimport { Embeddings, type EmbeddingsParams } from \"@langchain/core/embeddings\";\n\nimport { encodeApiKey } from \"../utils/zhipuai.js\";\n\n/**\n * Interface that extends EmbeddingsParams and defines additional\n * parameters specific to the ZhipuAIEmbeddingsParams class.\n */\nexport interface ZhipuAIEmbeddingsParams extends EmbeddingsParams {\n /**\n * Model Name to use\n */\n modelName?: \"embedding-2\" | \"embedding-3\";\n /**\n * ZhipuAI API key to use\n */\n apiKey?: string;\n /**\n * Whether to strip new lines from the input text.\n */\n stripNewLines?: boolean;\n}\n\ninterface EmbeddingData {\n embedding: number[];\n index: number;\n object: string;\n}\n\ninterface TokenUsage {\n completion_tokens: number;\n prompt_tokens: number;\n total_tokens: number;\n}\n\nexport interface ZhipuAIEmbeddingsResult {\n model: string;\n data: EmbeddingData[];\n object: string;\n usage: TokenUsage;\n}\n\nexport class ZhipuAIEmbeddings\n extends Embeddings\n implements ZhipuAIEmbeddingsParams\n{\n modelName: ZhipuAIEmbeddingsParams[\"modelName\"] = \"embedding-2\";\n\n apiKey?: string;\n\n stripNewLines = true;\n\n private embeddingsAPIURL = \"https://open.bigmodel.cn/api/paas/v4/embeddings\";\n\n constructor(fields?: ZhipuAIEmbeddingsParams) {\n super(fields ?? {});\n\n this.modelName = fields?.modelName ?? this.modelName;\n this.stripNewLines = fields?.stripNewLines ?? this.stripNewLines;\n this.apiKey = fields?.apiKey ?? getEnvironmentVariable(\"ZHIPUAI_API_KEY\");\n\n if (!this.apiKey) {\n throw new Error(\"ZhipuAI API key not found\");\n }\n }\n\n /**\n * Private method to make a request to the TogetherAI API to generate\n * embeddings. Handles the retry logic and returns the response from the API.\n * @param {string} input The input text to embed.\n * @returns Promise that resolves to the response from the API.\n * @TODO Figure out return type and statically type it.\n */\n private async embeddingWithRetry(\n input: string\n ): Promise<ZhipuAIEmbeddingsResult> {\n const text = this.stripNewLines ? input.replace(/\\n/g, \" \") : input;\n\n const body = JSON.stringify({ input: text, model: this.modelName });\n const headers = {\n Accept: \"application/json\",\n \"Content-Type\": \"application/json\",\n Authorization: encodeApiKey(this.apiKey),\n };\n\n return this.caller.call(async () => {\n const fetchResponse = await fetch(this.embeddingsAPIURL, {\n method: \"POST\",\n headers,\n body,\n });\n\n if (fetchResponse.status === 200) {\n return fetchResponse.json();\n }\n throw new Error(\n `Error getting embeddings from ZhipuAI. ${JSON.stringify(\n await fetchResponse.json(),\n null,\n 2\n )}`\n );\n });\n }\n\n /**\n * Method to generate an embedding for a single document. Calls the\n * embeddingWithRetry method with the document as the input.\n * @param {string} text Document to generate an embedding for.\n * @returns {Promise<number[]>} Promise that resolves to an embedding for the document.\n */\n async embedQuery(text: string): Promise<number[]> {\n const { data } = await this.embeddingWithRetry(text);\n return data[0].embedding;\n }\n\n /**\n * Method that takes an array of documents as input and returns a promise\n * that resolves to a 2D array of embeddings for each document. It calls\n * the embedQuery method for each document in the array.\n * @param documents Array of documents for which to generate embeddings.\n * @returns Promise that resolves to a 2D array of embeddings for each input document.\n */\n embedDocuments(documents: string[]): Promise<number[][]> {\n return Promise.all(documents.map((doc) => this.embedQuery(doc)));\n }\n}\n"],"mappings":";;;;;;;AA4CA,IAAa,oBAAb,cACUA,2BAAAA,WAEV;CACE,YAAkD;CAElD;CAEA,gBAAgB;CAEhB,mBAA2B;CAE3B,YAAY,QAAkC;AAC5C,QAAM,UAAU,EAAE,CAAC;AAEnB,OAAK,YAAY,QAAQ,aAAa,KAAK;AAC3C,OAAK,gBAAgB,QAAQ,iBAAiB,KAAK;AACnD,OAAK,SAAS,QAAQ,WAAA,GAAA,0BAAA,wBAAiC,kBAAkB;AAEzE,MAAI,CAAC,KAAK,OACR,OAAM,IAAI,MAAM,4BAA4B;;;;;;;;;CAWhD,MAAc,mBACZ,OACkC;EAClC,MAAM,OAAO,KAAK,gBAAgB,MAAM,QAAQ,OAAO,IAAI,GAAG;EAE9D,MAAM,OAAO,KAAK,UAAU;GAAE,OAAO;GAAM,OAAO,KAAK;GAAW,CAAC;EACnE,MAAM,UAAU;GACd,QAAQ;GACR,gBAAgB;GAChB,eAAeC,gBAAAA,aAAa,KAAK,OAAO;GACzC;AAED,SAAO,KAAK,OAAO,KAAK,YAAY;GAClC,MAAM,gBAAgB,MAAM,MAAM,KAAK,kBAAkB;IACvD,QAAQ;IACR;IACA;IACD,CAAC;AAEF,OAAI,cAAc,WAAW,IAC3B,QAAO,cAAc,MAAM;AAE7B,SAAM,IAAI,MACR,0CAA0C,KAAK,UAC7C,MAAM,cAAc,MAAM,EAC1B,MACA,EACD,GACF;IACD;;;;;;;;CASJ,MAAM,WAAW,MAAiC;EAChD,MAAM,EAAE,SAAS,MAAM,KAAK,mBAAmB,KAAK;AACpD,SAAO,KAAK,GAAG;;;;;;;;;CAUjB,eAAe,WAA0C;AACvD,SAAO,QAAQ,IAAI,UAAU,KAAK,QAAQ,KAAK,WAAW,IAAI,CAAC,CAAC"}