Spaces:
Running
Running
File size: 4,866 Bytes
c2b7eb3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | import { wrapOpenAIClientError } from "./utils/client.js";
import { getEndpoint, getHeadersWithUserAgent } from "./utils/azure.js";
import { OpenAI } from "openai";
import { getEnvironmentVariable } from "@langchain/core/utils/env";
import { chunkArray } from "@langchain/core/utils/chunk_array";
import { Embeddings } from "@langchain/core/embeddings";
//#region src/embeddings.ts
/**
* Class for generating embeddings using the OpenAI API.
*
* To use with Azure, import the `AzureOpenAIEmbeddings` class.
*
* @example
* ```typescript
* // Embed a query using OpenAIEmbeddings to generate embeddings for a given text
* const model = new OpenAIEmbeddings();
* const res = await model.embedQuery(
* "What would be a good company name for a company that makes colorful socks?",
* );
* console.log({ res });
*
* ```
*/
var OpenAIEmbeddings = class extends Embeddings {
model = "text-embedding-ada-002";
/** @deprecated Use "model" instead */
modelName;
batchSize = 512;
stripNewLines = true;
/**
* The number of dimensions the resulting output embeddings should have.
* Only supported in `text-embedding-3` and later models.
*/
dimensions;
timeout;
organization;
encodingFormat;
client;
clientConfig;
apiKey;
constructor(fields) {
const fieldsWithDefaults = {
maxConcurrency: 2,
...fields
};
super(fieldsWithDefaults);
const apiKey = fieldsWithDefaults?.apiKey ?? fieldsWithDefaults?.openAIApiKey ?? getEnvironmentVariable("OPENAI_API_KEY");
this.organization = fieldsWithDefaults?.configuration?.organization ?? getEnvironmentVariable("OPENAI_ORGANIZATION");
this.model = fieldsWithDefaults?.model ?? fieldsWithDefaults?.modelName ?? this.model;
this.modelName = this.model;
this.batchSize = fieldsWithDefaults?.batchSize ?? this.batchSize;
this.stripNewLines = fieldsWithDefaults?.stripNewLines ?? this.stripNewLines;
this.timeout = fieldsWithDefaults?.timeout;
this.dimensions = fieldsWithDefaults?.dimensions;
this.encodingFormat = fieldsWithDefaults?.encodingFormat;
this.clientConfig = {
apiKey,
organization: this.organization,
dangerouslyAllowBrowser: true,
...fields?.configuration
};
}
/**
* Method to generate embeddings for an array of documents. Splits the
* documents into batches and makes requests to the OpenAI API to generate
* embeddings.
* @param texts Array of documents to generate embeddings for.
* @returns Promise that resolves to a 2D array of embeddings for each document.
*/
async embedDocuments(texts) {
const batches = chunkArray(this.stripNewLines ? texts.map((t) => t.replace(/\n/g, " ")) : texts, this.batchSize);
const batchRequests = batches.map((batch) => {
const params = {
model: this.model,
input: batch
};
if (this.dimensions) params.dimensions = this.dimensions;
if (this.encodingFormat) params.encoding_format = this.encodingFormat;
return this.embeddingWithRetry(params);
});
const batchResponses = await Promise.all(batchRequests);
const embeddings = [];
for (let i = 0; i < batchResponses.length; i += 1) {
const batch = batches[i];
const { data: batchResponse } = batchResponses[i];
for (let j = 0; j < batch.length; j += 1) embeddings.push(batchResponse[j].embedding);
}
return embeddings;
}
/**
* Method to generate an embedding for a single document. Calls the
* embeddingWithRetry method with the document as the input.
* @param text Document to generate an embedding for.
* @returns Promise that resolves to an embedding for the document.
*/
async embedQuery(text) {
const params = {
model: this.model,
input: this.stripNewLines ? text.replace(/\n/g, " ") : text
};
if (this.dimensions) params.dimensions = this.dimensions;
if (this.encodingFormat) params.encoding_format = this.encodingFormat;
const { data } = await this.embeddingWithRetry(params);
return data[0].embedding;
}
/**
* Private method to make a request to the OpenAI API to generate
* embeddings. Handles the retry logic and returns the response from the
* API.
* @param request Request to send to the OpenAI API.
* @returns Promise that resolves to the response from the API.
*/
async embeddingWithRetry(request) {
if (!this.client) {
const endpoint = getEndpoint({ baseURL: this.clientConfig.baseURL });
const params = {
...this.clientConfig,
baseURL: endpoint,
timeout: this.timeout,
maxRetries: 0
};
if (!params.baseURL) delete params.baseURL;
params.defaultHeaders = getHeadersWithUserAgent(params.defaultHeaders);
this.client = new OpenAI(params);
}
const requestOptions = {};
return this.caller.call(async () => {
try {
return await this.client.embeddings.create(request, requestOptions);
} catch (e) {
throw wrapOpenAIClientError(e);
}
});
}
};
//#endregion
export { OpenAIEmbeddings };
//# sourceMappingURL=embeddings.js.map |