78 lines
2.8 KiB
JavaScript
78 lines
2.8 KiB
JavaScript
|
"use strict";
|
||
|
|
Object.defineProperty(exports, "__esModule", { value: true });
|
||
|
|
exports.CloudflareWorkersAIEmbeddings = void 0;
|
||
|
|
const ai_1 = require("@cloudflare/ai");
|
||
|
|
const embeddings_1 = require("@langchain/core/embeddings");
|
||
|
|
const chunk_array_1 = require("@langchain/core/utils/chunk_array");
|
||
|
|
/** @deprecated Install and import from "@langchain/cloudflare" instead. */
|
||
|
|
class CloudflareWorkersAIEmbeddings extends embeddings_1.Embeddings {
|
||
|
|
constructor(fields) {
|
||
|
|
super(fields);
|
||
|
|
Object.defineProperty(this, "modelName", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: "@cf/baai/bge-base-en-v1.5"
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "model", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: "@cf/baai/bge-base-en-v1.5"
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "batchSize", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: 50
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "stripNewLines", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: true
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "ai", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
if (!fields.binding) {
|
||
|
|
throw new Error("Must supply a Workers AI binding, eg { binding: env.AI }");
|
||
|
|
}
|
||
|
|
this.ai = new ai_1.Ai(fields.binding);
|
||
|
|
this.modelName = fields?.model ?? fields.modelName ?? this.model;
|
||
|
|
this.model = this.modelName;
|
||
|
|
this.stripNewLines = fields.stripNewLines ?? this.stripNewLines;
|
||
|
|
}
|
||
|
|
async embedDocuments(texts) {
|
||
|
|
const batches = (0, chunk_array_1.chunkArray)(this.stripNewLines ? texts.map((t) => t.replace(/\n/g, " ")) : texts, this.batchSize);
|
||
|
|
const batchRequests = batches.map((batch) => this.runEmbedding(batch));
|
||
|
|
const batchResponses = await Promise.all(batchRequests);
|
||
|
|
const embeddings = [];
|
||
|
|
for (let i = 0; i < batchResponses.length; i += 1) {
|
||
|
|
const batchResponse = batchResponses[i];
|
||
|
|
for (let j = 0; j < batchResponse.length; j += 1) {
|
||
|
|
embeddings.push(batchResponse[j]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return embeddings;
|
||
|
|
}
|
||
|
|
async embedQuery(text) {
|
||
|
|
const data = await this.runEmbedding([
|
||
|
|
this.stripNewLines ? text.replace(/\n/g, " ") : text,
|
||
|
|
]);
|
||
|
|
return data[0];
|
||
|
|
}
|
||
|
|
async runEmbedding(texts) {
|
||
|
|
return this.caller.call(async () => {
|
||
|
|
const response = await this.ai.run(this.model, {
|
||
|
|
text: texts,
|
||
|
|
});
|
||
|
|
return response.data;
|
||
|
|
});
|
||
|
|
}
|
||
|
|
}
|
||
|
|
exports.CloudflareWorkersAIEmbeddings = CloudflareWorkersAIEmbeddings;
|