-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
typeorm.ts
301 lines (266 loc) Β· 9.36 KB
/
typeorm.ts
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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
import { Metadata } from "@opensearch-project/opensearch/api/types.js";
import { DataSource, DataSourceOptions, EntitySchema } from "typeorm";
import type { EmbeddingsInterface } from "@langchain/core/embeddings";
import { VectorStore } from "@langchain/core/vectorstores";
import { Document } from "@langchain/core/documents";
import { getEnvironmentVariable } from "@langchain/core/utils/env";
/**
* Interface that defines the arguments required to create a
* `TypeORMVectorStore` instance. It includes Postgres connection options,
* table name, filter, and verbosity level.
*/
export interface TypeORMVectorStoreArgs {
postgresConnectionOptions: DataSourceOptions;
tableName?: string;
filter?: Metadata;
verbose?: boolean;
}
/**
* Class that extends the `Document` base class and adds an `embedding`
* property. It represents a document in the vector store.
*/
export class TypeORMVectorStoreDocument extends Document {
embedding: string;
id?: string;
}
const defaultDocumentTableName = "documents";
/**
* Class that provides an interface to a Postgres vector database. It
* extends the `VectorStore` base class and implements methods for adding
* documents and vectors, performing similarity searches, and ensuring the
* existence of a table in the database.
*/
export class TypeORMVectorStore extends VectorStore {
declare FilterType: Metadata;
tableName: string;
documentEntity: EntitySchema;
filter?: Metadata;
appDataSource: DataSource;
_verbose?: boolean;
_vectorstoreType(): string {
return "typeorm";
}
private constructor(
embeddings: EmbeddingsInterface,
fields: TypeORMVectorStoreArgs
) {
super(embeddings, fields);
this.tableName = fields.tableName || defaultDocumentTableName;
this.filter = fields.filter;
const TypeORMDocumentEntity = new EntitySchema<TypeORMVectorStoreDocument>({
name: fields.tableName ?? defaultDocumentTableName,
columns: {
id: {
generated: "uuid",
type: "uuid",
primary: true,
},
pageContent: {
type: String,
},
metadata: {
type: "jsonb",
},
embedding: {
type: String,
},
},
});
const appDataSource = new DataSource({
entities: [TypeORMDocumentEntity],
...fields.postgresConnectionOptions,
});
this.appDataSource = appDataSource;
this.documentEntity = TypeORMDocumentEntity;
this._verbose =
getEnvironmentVariable("LANGCHAIN_VERBOSE") === "true" ??
fields.verbose ??
false;
}
/**
* Static method to create a new `TypeORMVectorStore` instance from a
* `DataSource`. It initializes the `DataSource` if it is not already
* initialized.
* @param embeddings Embeddings instance.
* @param fields `TypeORMVectorStoreArgs` instance.
* @returns A new instance of `TypeORMVectorStore`.
*/
static async fromDataSource(
embeddings: EmbeddingsInterface,
fields: TypeORMVectorStoreArgs
): Promise<TypeORMVectorStore> {
const postgresqlVectorStore = new TypeORMVectorStore(embeddings, fields);
if (!postgresqlVectorStore.appDataSource.isInitialized) {
await postgresqlVectorStore.appDataSource.initialize();
}
return postgresqlVectorStore;
}
/**
* Method to add documents to the vector store. It ensures the existence
* of the table in the database, converts the documents into vectors, and
* adds them to the store.
* @param documents Array of `Document` instances.
* @returns Promise that resolves when the documents have been added.
*/
async addDocuments(documents: Document[]): Promise<void> {
const texts = documents.map(({ pageContent }) => pageContent);
// This will create the table if it does not exist. We can call it every time as it doesn't
// do anything if the table already exists, and it is not expensive in terms of performance
await this.ensureTableInDatabase();
return this.addVectors(
await this.embeddings.embedDocuments(texts),
documents
);
}
/**
* Method to add vectors to the vector store. It converts the vectors into
* rows and inserts them into the database.
* @param vectors Array of vectors.
* @param documents Array of `Document` instances.
* @returns Promise that resolves when the vectors have been added.
*/
async addVectors(vectors: number[][], documents: Document[]): Promise<void> {
const rows = vectors.map((embedding, idx) => {
const embeddingString = `[${embedding.join(",")}]`;
const documentRow = {
pageContent: documents[idx].pageContent,
embedding: embeddingString,
metadata: documents[idx].metadata,
};
return documentRow;
});
const documentRepository = this.appDataSource.getRepository(
this.documentEntity
);
const chunkSize = 500;
for (let i = 0; i < rows.length; i += chunkSize) {
const chunk = rows.slice(i, i + chunkSize);
try {
await documentRepository.save(chunk);
} catch (e) {
console.error(e);
throw new Error(`Error inserting: ${chunk[0].pageContent}`);
}
}
}
/**
* Method to perform a similarity search in the vector store. It returns
* the `k` most similar documents to the query vector, along with their
* similarity scores.
* @param query Query vector.
* @param k Number of most similar documents to return.
* @param filter Optional filter to apply to the search.
* @returns Promise that resolves with an array of tuples, each containing a `TypeORMVectorStoreDocument` and its similarity score.
*/
async similaritySearchVectorWithScore(
query: number[],
k: number,
filter?: this["FilterType"]
): Promise<[TypeORMVectorStoreDocument, number][]> {
const embeddingString = `[${query.join(",")}]`;
const _filter = filter ?? "{}";
const queryString = `
SELECT *, embedding <=> $1 as "_distance"
FROM ${this.tableName}
WHERE metadata @> $2
ORDER BY "_distance" ASC
LIMIT $3;`;
const documents = await this.appDataSource.query(queryString, [
embeddingString,
_filter,
k,
]);
const results = [] as [TypeORMVectorStoreDocument, number][];
for (const doc of documents) {
if (doc._distance != null && doc.pageContent != null) {
const document = new Document(doc) as TypeORMVectorStoreDocument;
document.id = doc.id;
results.push([document, doc._distance]);
}
}
return results;
}
/**
* Method to ensure the existence of the table in the database. It creates
* the table if it does not already exist.
* @returns Promise that resolves when the table has been ensured.
*/
async ensureTableInDatabase(): Promise<void> {
await this.appDataSource.query("CREATE EXTENSION IF NOT EXISTS vector;");
await this.appDataSource.query(
'CREATE EXTENSION IF NOT EXISTS "uuid-ossp";'
);
await this.appDataSource.query(`
CREATE TABLE IF NOT EXISTS ${this.tableName} (
"id" uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
"pageContent" text,
metadata jsonb,
embedding vector
);
`);
}
/**
* Static method to create a new `TypeORMVectorStore` instance from an
* array of texts and their metadata. It converts the texts into
* `Document` instances and adds them to the store.
* @param texts Array of texts.
* @param metadatas Array of metadata objects or a single metadata object.
* @param embeddings Embeddings instance.
* @param dbConfig `TypeORMVectorStoreArgs` instance.
* @returns Promise that resolves with a new instance of `TypeORMVectorStore`.
*/
static async fromTexts(
texts: string[],
metadatas: object[] | object,
embeddings: EmbeddingsInterface,
dbConfig: TypeORMVectorStoreArgs
): Promise<TypeORMVectorStore> {
const docs = [];
for (let i = 0; i < texts.length; i += 1) {
const metadata = Array.isArray(metadatas) ? metadatas[i] : metadatas;
const newDoc = new Document({
pageContent: texts[i],
metadata,
});
docs.push(newDoc);
}
return TypeORMVectorStore.fromDocuments(docs, embeddings, dbConfig);
}
/**
* Static method to create a new `TypeORMVectorStore` instance from an
* array of `Document` instances. It adds the documents to the store.
* @param docs Array of `Document` instances.
* @param embeddings Embeddings instance.
* @param dbConfig `TypeORMVectorStoreArgs` instance.
* @returns Promise that resolves with a new instance of `TypeORMVectorStore`.
*/
static async fromDocuments(
docs: Document[],
embeddings: EmbeddingsInterface,
dbConfig: TypeORMVectorStoreArgs
): Promise<TypeORMVectorStore> {
const instance = await TypeORMVectorStore.fromDataSource(
embeddings,
dbConfig
);
await instance.addDocuments(docs);
return instance;
}
/**
* Static method to create a new `TypeORMVectorStore` instance from an
* existing index.
* @param embeddings Embeddings instance.
* @param dbConfig `TypeORMVectorStoreArgs` instance.
* @returns Promise that resolves with a new instance of `TypeORMVectorStore`.
*/
static async fromExistingIndex(
embeddings: EmbeddingsInterface,
dbConfig: TypeORMVectorStoreArgs
): Promise<TypeORMVectorStore> {
const instance = await TypeORMVectorStore.fromDataSource(
embeddings,
dbConfig
);
return instance;
}
}