305 lines
12 KiB
JavaScript
305 lines
12 KiB
JavaScript
|
"use strict";
|
||
|
|
Object.defineProperty(exports, "__esModule", { value: true });
|
||
|
|
exports.PrismaVectorStore = void 0;
|
||
|
|
const documents_1 = require("@langchain/core/documents");
|
||
|
|
const vectorstores_1 = require("@langchain/core/vectorstores");
|
||
|
|
const IdColumnSymbol = Symbol("id");
|
||
|
|
const ContentColumnSymbol = Symbol("content");
|
||
|
|
const OpMap = {
|
||
|
|
equals: "=",
|
||
|
|
in: "IN",
|
||
|
|
notIn: "NOT IN",
|
||
|
|
isNull: "IS NULL",
|
||
|
|
isNotNull: "IS NOT NULL",
|
||
|
|
like: "LIKE",
|
||
|
|
lt: "<",
|
||
|
|
lte: "<=",
|
||
|
|
gt: ">",
|
||
|
|
gte: ">=",
|
||
|
|
not: "<>",
|
||
|
|
};
|
||
|
|
/**
|
||
|
|
* A specific implementation of the VectorStore class that is designed to
|
||
|
|
* work with Prisma. It provides methods for adding models, documents, and
|
||
|
|
* vectors, as well as for performing similarity searches.
|
||
|
|
*/
|
||
|
|
class PrismaVectorStore extends vectorstores_1.VectorStore {
|
||
|
|
_vectorstoreType() {
|
||
|
|
return "prisma";
|
||
|
|
}
|
||
|
|
constructor(embeddings, config) {
|
||
|
|
super(embeddings, {});
|
||
|
|
Object.defineProperty(this, "tableName", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "vectorColumnName", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "selectColumns", {
|
||
|
|
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, "idColumn", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "contentColumn", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "db", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
Object.defineProperty(this, "Prisma", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: void 0
|
||
|
|
});
|
||
|
|
this.Prisma = config.prisma;
|
||
|
|
this.db = config.db;
|
||
|
|
const entries = Object.entries(config.columns);
|
||
|
|
const idColumn = entries.find((i) => i[1] === IdColumnSymbol)?.[0];
|
||
|
|
const contentColumn = entries.find((i) => i[1] === ContentColumnSymbol)?.[0];
|
||
|
|
if (idColumn == null)
|
||
|
|
throw new Error("Missing ID column");
|
||
|
|
if (contentColumn == null)
|
||
|
|
throw new Error("Missing content column");
|
||
|
|
this.idColumn = idColumn;
|
||
|
|
this.contentColumn = contentColumn;
|
||
|
|
this.tableName = config.tableName;
|
||
|
|
this.vectorColumnName = config.vectorColumnName;
|
||
|
|
this.selectColumns = entries
|
||
|
|
.map(([key, alias]) => (alias && key) || null)
|
||
|
|
.filter((x) => !!x);
|
||
|
|
if (config.filter) {
|
||
|
|
this.filter = config.filter;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Creates a new PrismaVectorStore with the specified model.
|
||
|
|
* @param db The PrismaClient instance.
|
||
|
|
* @returns An object with create, fromTexts, and fromDocuments methods.
|
||
|
|
*/
|
||
|
|
static withModel(db) {
|
||
|
|
function create(embeddings, config) {
|
||
|
|
return new PrismaVectorStore(embeddings, { ...config, db });
|
||
|
|
}
|
||
|
|
async function 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 newDoc = new documents_1.Document({
|
||
|
|
pageContent: texts[i],
|
||
|
|
metadata,
|
||
|
|
});
|
||
|
|
docs.push(newDoc);
|
||
|
|
}
|
||
|
|
return PrismaVectorStore.fromDocuments(docs, embeddings, {
|
||
|
|
...dbConfig,
|
||
|
|
db,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
async function fromDocuments(docs, embeddings, dbConfig) {
|
||
|
|
const instance = new PrismaVectorStore(embeddings, { ...dbConfig, db });
|
||
|
|
await instance.addDocuments(docs);
|
||
|
|
return instance;
|
||
|
|
}
|
||
|
|
return { create, fromTexts, fromDocuments };
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Adds the specified models to the store.
|
||
|
|
* @param models The models to add.
|
||
|
|
* @returns A promise that resolves when the models have been added.
|
||
|
|
*/
|
||
|
|
async addModels(models) {
|
||
|
|
return this.addDocuments(models.map((metadata) => {
|
||
|
|
const pageContent = metadata[this.contentColumn];
|
||
|
|
if (typeof pageContent !== "string")
|
||
|
|
throw new Error("Content column must be a string");
|
||
|
|
return new documents_1.Document({ pageContent, metadata });
|
||
|
|
}));
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Adds the specified documents to the store.
|
||
|
|
* @param documents The documents to add.
|
||
|
|
* @returns A promise that resolves when the documents have been added.
|
||
|
|
*/
|
||
|
|
async addDocuments(documents) {
|
||
|
|
const texts = documents.map(({ pageContent }) => pageContent);
|
||
|
|
return this.addVectors(await this.embeddings.embedDocuments(texts), documents);
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Adds the specified vectors to the store.
|
||
|
|
* @param vectors The vectors to add.
|
||
|
|
* @param documents The documents associated with the vectors.
|
||
|
|
* @returns A promise that resolves when the vectors have been added.
|
||
|
|
*/
|
||
|
|
async addVectors(vectors, documents) {
|
||
|
|
// table name, column name cannot be parametrised
|
||
|
|
// these fields are thus not escaped by Prisma and can be dangerous if user input is used
|
||
|
|
const idColumnRaw = this.Prisma.raw(`"${this.idColumn}"`);
|
||
|
|
const tableNameRaw = this.Prisma.raw(`"${this.tableName}"`);
|
||
|
|
const vectorColumnRaw = this.Prisma.raw(`"${this.vectorColumnName}"`);
|
||
|
|
await this.db.$transaction(vectors.map((vector, idx) => this.db.$executeRaw(this.Prisma.sql `UPDATE ${tableNameRaw}
|
||
|
|
SET ${vectorColumnRaw} = ${`[${vector.join(",")}]`}::vector
|
||
|
|
WHERE ${idColumnRaw} = ${documents[idx].metadata[this.idColumn]}
|
||
|
|
`)));
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Performs a similarity search with the specified query.
|
||
|
|
* @param query The query to use for the similarity search.
|
||
|
|
* @param k The number of results to return.
|
||
|
|
* @param _filter The filter to apply to the results.
|
||
|
|
* @param _callbacks The callbacks to use during the search.
|
||
|
|
* @returns A promise that resolves with the search results.
|
||
|
|
*/
|
||
|
|
async similaritySearch(query, k = 4, filter = undefined) {
|
||
|
|
const results = await this.similaritySearchVectorWithScore(await this.embeddings.embedQuery(query), k, filter);
|
||
|
|
return results.map((result) => result[0]);
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Performs a similarity search with the specified query and returns the
|
||
|
|
* results along with their scores.
|
||
|
|
* @param query The query to use for the similarity search.
|
||
|
|
* @param k The number of results to return.
|
||
|
|
* @param filter The filter to apply to the results.
|
||
|
|
* @param _callbacks The callbacks to use during the search.
|
||
|
|
* @returns A promise that resolves with the search results and their scores.
|
||
|
|
*/
|
||
|
|
async similaritySearchWithScore(query, k, filter) {
|
||
|
|
return super.similaritySearchWithScore(query, k, filter);
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Performs a similarity search with the specified vector and returns the
|
||
|
|
* results along with their scores.
|
||
|
|
* @param query The vector to use for the similarity search.
|
||
|
|
* @param k The number of results to return.
|
||
|
|
* @param filter The filter to apply to the results.
|
||
|
|
* @returns A promise that resolves with the search results and their scores.
|
||
|
|
*/
|
||
|
|
async similaritySearchVectorWithScore(query, k, filter) {
|
||
|
|
// table name, column names cannot be parametrised
|
||
|
|
// these fields are thus not escaped by Prisma and can be dangerous if user input is used
|
||
|
|
const vectorColumnRaw = this.Prisma.raw(`"${this.vectorColumnName}"`);
|
||
|
|
const tableNameRaw = this.Prisma.raw(`"${this.tableName}"`);
|
||
|
|
const selectRaw = this.Prisma.raw(this.selectColumns.map((x) => `"${x}"`).join(", "));
|
||
|
|
const vector = `[${query.join(",")}]`;
|
||
|
|
const articles = await this.db.$queryRaw(this.Prisma.join([
|
||
|
|
this.Prisma.sql `
|
||
|
|
SELECT ${selectRaw}, ${vectorColumnRaw} <=> ${vector}::vector as "_distance"
|
||
|
|
FROM ${tableNameRaw}
|
||
|
|
`,
|
||
|
|
this.buildSqlFilterStr(filter ?? this.filter),
|
||
|
|
this.Prisma.sql `
|
||
|
|
ORDER BY "_distance" ASC
|
||
|
|
LIMIT ${k};
|
||
|
|
`,
|
||
|
|
].filter((x) => x != null), ""));
|
||
|
|
const results = [];
|
||
|
|
for (const article of articles) {
|
||
|
|
if (article._distance != null && article[this.contentColumn] != null) {
|
||
|
|
results.push([
|
||
|
|
new documents_1.Document({
|
||
|
|
pageContent: article[this.contentColumn],
|
||
|
|
metadata: article,
|
||
|
|
}),
|
||
|
|
article._distance,
|
||
|
|
]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return results;
|
||
|
|
}
|
||
|
|
buildSqlFilterStr(filter) {
|
||
|
|
if (filter == null)
|
||
|
|
return null;
|
||
|
|
return this.Prisma.join(Object.entries(filter).flatMap(([key, ops]) => Object.entries(ops).map(([opName, value]) => {
|
||
|
|
// column name, operators cannot be parametrised
|
||
|
|
// these fields are thus not escaped by Prisma and can be dangerous if user input is used
|
||
|
|
const opNameKey = opName;
|
||
|
|
const colRaw = this.Prisma.raw(`"${key}"`);
|
||
|
|
const opRaw = this.Prisma.raw(OpMap[opNameKey]);
|
||
|
|
switch (OpMap[opNameKey]) {
|
||
|
|
case OpMap.notIn:
|
||
|
|
case OpMap.in: {
|
||
|
|
if (!Array.isArray(value)) {
|
||
|
|
throw new Error(`Invalid filter: IN operator requires an array. Received: ${JSON.stringify(value, null, 2)}`);
|
||
|
|
}
|
||
|
|
return this.Prisma.sql `${colRaw} ${opRaw} (${this.Prisma.join(value)})`;
|
||
|
|
}
|
||
|
|
case OpMap.isNull:
|
||
|
|
case OpMap.isNotNull:
|
||
|
|
return this.Prisma.sql `${colRaw} ${opRaw}`;
|
||
|
|
default:
|
||
|
|
return this.Prisma.sql `${colRaw} ${opRaw} ${value}`;
|
||
|
|
}
|
||
|
|
})), " AND ", " WHERE ");
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Creates a new PrismaVectorStore from the specified texts.
|
||
|
|
* @param texts The texts to use to create the store.
|
||
|
|
* @param metadatas The metadata for the texts.
|
||
|
|
* @param embeddings The embeddings to use.
|
||
|
|
* @param dbConfig The database configuration.
|
||
|
|
* @returns A promise that resolves with the new PrismaVectorStore.
|
||
|
|
*/
|
||
|
|
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 newDoc = new documents_1.Document({
|
||
|
|
pageContent: texts[i],
|
||
|
|
metadata,
|
||
|
|
});
|
||
|
|
docs.push(newDoc);
|
||
|
|
}
|
||
|
|
return PrismaVectorStore.fromDocuments(docs, embeddings, dbConfig);
|
||
|
|
}
|
||
|
|
/**
|
||
|
|
* Creates a new PrismaVectorStore from the specified documents.
|
||
|
|
* @param docs The documents to use to create the store.
|
||
|
|
* @param embeddings The embeddings to use.
|
||
|
|
* @param dbConfig The database configuration.
|
||
|
|
* @returns A promise that resolves with the new PrismaVectorStore.
|
||
|
|
*/
|
||
|
|
static async fromDocuments(docs, embeddings, dbConfig) {
|
||
|
|
const instance = new PrismaVectorStore(embeddings, dbConfig);
|
||
|
|
await instance.addDocuments(docs);
|
||
|
|
return instance;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
exports.PrismaVectorStore = PrismaVectorStore;
|
||
|
|
Object.defineProperty(PrismaVectorStore, "IdColumn", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: IdColumnSymbol
|
||
|
|
});
|
||
|
|
Object.defineProperty(PrismaVectorStore, "ContentColumn", {
|
||
|
|
enumerable: true,
|
||
|
|
configurable: true,
|
||
|
|
writable: true,
|
||
|
|
value: ContentColumnSymbol
|
||
|
|
});
|