234 lines
9.8 KiB
JavaScript
234 lines
9.8 KiB
JavaScript
import * as uuid from "uuid";
|
|
import { VectorStore } from "@langchain/core/vectorstores";
|
|
import { Document } from "@langchain/core/documents";
|
|
import { chunkArray } from "@langchain/core/utils/chunk_array";
|
|
import { FakeEmbeddings } from "@langchain/core/utils/testing";
|
|
import { AsyncCaller, } from "@langchain/core/utils/async_caller";
|
|
const CONCURRENT_UPSERT_LIMIT = 1000;
|
|
/**
|
|
* The main class that extends the 'VectorStore' class. It provides
|
|
* methods for interacting with Upstash index, such as adding documents,
|
|
* deleting documents, performing similarity search and more.
|
|
*/
|
|
export class UpstashVectorStore extends VectorStore {
|
|
_vectorstoreType() {
|
|
return "upstash";
|
|
}
|
|
constructor(embeddings, args) {
|
|
super(embeddings, args);
|
|
Object.defineProperty(this, "index", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "caller", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "useUpstashEmbeddings", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "filter", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
Object.defineProperty(this, "namespace", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
// Special case where the embeddings instance is a FakeEmbeddings instance. In this case, we need to disable "instanceof" rule.
|
|
// eslint-disable-next-line no-instanceof/no-instanceof
|
|
if (embeddings instanceof FakeEmbeddings) {
|
|
this.useUpstashEmbeddings = true;
|
|
}
|
|
const { index, namespace, ...asyncCallerArgs } = args;
|
|
this.index = index;
|
|
this.caller = new AsyncCaller(asyncCallerArgs);
|
|
this.filter = args.filter;
|
|
this.namespace = namespace;
|
|
}
|
|
/**
|
|
* This method adds documents to Upstash database. Documents are first converted to vectors
|
|
* using the provided embeddings instance, and then upserted to the database.
|
|
* @param documents Array of Document objects to be added to the database.
|
|
* @param options Optional object containing array of ids for the documents.
|
|
* @returns Promise that resolves with the ids of the provided documents when the upsert operation is done.
|
|
*/
|
|
async addDocuments(documents, options) {
|
|
const texts = documents.map(({ pageContent }) => pageContent);
|
|
if (this.useUpstashEmbeddings || options?.useUpstashEmbeddings) {
|
|
return this._addData(documents, options);
|
|
}
|
|
const embeddings = await this.embeddings.embedDocuments(texts);
|
|
return this.addVectors(embeddings, documents, options);
|
|
}
|
|
/**
|
|
* This method adds the provided vectors to Upstash database.
|
|
* @param vectors Array of vectors to be added to the Upstash database.
|
|
* @param documents Array of Document objects, each associated with a vector.
|
|
* @param options Optional object containing the array of ids foor the vectors.
|
|
* @returns Promise that resolves with the ids of the provided documents when the upsert operation is done.
|
|
*/
|
|
async addVectors(vectors, documents, options) {
|
|
const documentIds = options?.ids ?? Array.from({ length: vectors.length }, () => uuid.v4());
|
|
const upstashVectors = vectors.map((vector, index) => {
|
|
const metadata = {
|
|
_pageContentLC: documents[index].pageContent,
|
|
...documents[index].metadata,
|
|
};
|
|
const id = documentIds[index];
|
|
return {
|
|
id,
|
|
vector,
|
|
metadata,
|
|
};
|
|
});
|
|
const namespace = this.index.namespace(this.namespace ?? "");
|
|
const vectorChunks = chunkArray(upstashVectors, CONCURRENT_UPSERT_LIMIT);
|
|
const batchRequests = vectorChunks.map((chunk) => this.caller.call(async () => namespace.upsert(chunk)));
|
|
await Promise.all(batchRequests);
|
|
return documentIds;
|
|
}
|
|
/**
|
|
* This method adds the provided documents to Upstash database. The pageContent of the documents will be embedded by Upstash Embeddings.
|
|
* @param documents Array of Document objects to be added to the Upstash database.
|
|
* @param options Optional object containing the array of ids for the documents.
|
|
* @returns Promise that resolves with the ids of the provided documents when the upsert operation is done.
|
|
*/
|
|
async _addData(documents, options) {
|
|
const documentIds = options?.ids ?? Array.from({ length: documents.length }, () => uuid.v4());
|
|
const upstashVectorsWithData = documents.map((document, index) => {
|
|
const metadata = {
|
|
_pageContentLC: documents[index].pageContent,
|
|
...documents[index].metadata,
|
|
};
|
|
const id = documentIds[index];
|
|
return {
|
|
id,
|
|
data: document.pageContent,
|
|
metadata,
|
|
};
|
|
});
|
|
const namespace = this.index.namespace(this.namespace ?? "");
|
|
const vectorChunks = chunkArray(upstashVectorsWithData, CONCURRENT_UPSERT_LIMIT);
|
|
const batchRequests = vectorChunks.map((chunk) => this.caller.call(async () => namespace.upsert(chunk)));
|
|
await Promise.all(batchRequests);
|
|
return documentIds;
|
|
}
|
|
/**
|
|
* This method deletes documents from the Upstash database. You can either
|
|
* provide the target ids, or delete all vectors in the database.
|
|
* @param params Object containing either array of ids of the documents or boolean deleteAll.
|
|
* @returns Promise that resolves when the specified documents have been deleted from the database.
|
|
*/
|
|
async delete(params) {
|
|
const namespace = this.index.namespace(this.namespace ?? "");
|
|
if (params.deleteAll) {
|
|
await namespace.reset();
|
|
}
|
|
else if (params.ids) {
|
|
await namespace.delete(params.ids);
|
|
}
|
|
}
|
|
async _runUpstashQuery(query, k, filter, options) {
|
|
let queryResult = [];
|
|
const namespace = this.index.namespace(this.namespace ?? "");
|
|
if (typeof query === "string") {
|
|
queryResult = await namespace.query({
|
|
data: query,
|
|
topK: k,
|
|
includeMetadata: true,
|
|
filter,
|
|
...options,
|
|
});
|
|
}
|
|
else {
|
|
queryResult = await namespace.query({
|
|
vector: query,
|
|
topK: k,
|
|
includeMetadata: true,
|
|
filter,
|
|
...options,
|
|
});
|
|
}
|
|
return queryResult;
|
|
}
|
|
/**
|
|
* This method performs a similarity search in the Upstash database
|
|
* over the existing vectors.
|
|
* @param query Query vector for the similarity search.
|
|
* @param k The number of similar vectors to return as result.
|
|
* @returns Promise that resolves with an array of tuples, each containing
|
|
* Document object and similarity score. The length of the result will be
|
|
* maximum of 'k' and vectors in the index.
|
|
*/
|
|
async similaritySearchVectorWithScore(query, k, filter) {
|
|
const results = await this._runUpstashQuery(query, k, filter);
|
|
const searchResult = results.map((res) => {
|
|
const { _pageContentLC, ...metadata } = (res.metadata ??
|
|
{});
|
|
return [
|
|
new Document({
|
|
metadata,
|
|
pageContent: _pageContentLC,
|
|
}),
|
|
res.score,
|
|
];
|
|
});
|
|
return searchResult;
|
|
}
|
|
/**
|
|
* This method creates a new UpstashVector instance from an array of texts.
|
|
* The texts are initially converted to Document instances and added to Upstash
|
|
* database.
|
|
* @param texts The texts to create the documents from.
|
|
* @param metadatas The metadata values associated with the texts.
|
|
* @param embeddings Embedding interface of choice, to create the text embeddings.
|
|
* @param dbConfig Object containing the Upstash database configs.
|
|
* @returns Promise that resolves with a new UpstashVector instance.
|
|
*/
|
|
static async fromTexts(texts, metadatas, embeddings, dbConfig) {
|
|
const docs = [];
|
|
for (let i = 0; i < texts.length; i += 1) {
|
|
const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;
|
|
const newDocument = new Document({
|
|
pageContent: texts[i],
|
|
metadata,
|
|
});
|
|
docs.push(newDocument);
|
|
}
|
|
return this.fromDocuments(docs, embeddings, dbConfig);
|
|
}
|
|
/**
|
|
* This method creates a new UpstashVector instance from an array of Document instances.
|
|
* @param docs The docs to be added to Upstash database.
|
|
* @param embeddings Embedding interface of choice, to create the embeddings.
|
|
* @param dbConfig Object containing the Upstash database configs.
|
|
* @returns Promise that resolves with a new UpstashVector instance
|
|
*/
|
|
static async fromDocuments(docs, embeddings, dbConfig) {
|
|
const instance = new this(embeddings, dbConfig);
|
|
await instance.addDocuments(docs);
|
|
return instance;
|
|
}
|
|
/**
|
|
* This method creates a new UpstashVector instance from an existing index.
|
|
* @param embeddings Embedding interface of the choice, to create the embeddings.
|
|
* @param dbConfig Object containing the Upstash database configs.
|
|
* @returns
|
|
*/
|
|
static async fromExistingIndex(embeddings, dbConfig) {
|
|
const instance = new this(embeddings, dbConfig);
|
|
return instance;
|
|
}
|
|
}
|