-
Notifications
You must be signed in to change notification settings - Fork 2k
/
indexing.ts
383 lines (331 loc) Β· 10.8 KB
/
indexing.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
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
import { VectorStore } from "@langchain/core/vectorstores";
import { v5 as uuidv5 } from "uuid";
import {
RecordManagerInterface,
UUIDV5_NAMESPACE,
} from "@langchain/community/indexes/base";
import { insecureHash } from "@langchain/core/utils/hash";
import { DocumentInterface } from "@langchain/core/documents";
import { BaseDocumentLoader } from "../document_loaders/base.js";
import { Document } from "../document.js";
type Metadata = Record<string, unknown>;
type IndexingResult = {
numAdded: number;
numDeleted: number;
numUpdated: number;
numSkipped: number;
};
type StringOrDocFunc = string | ((doc: DocumentInterface) => string);
export interface HashedDocumentInterface extends DocumentInterface {
uid: string;
hash_?: string;
contentHash?: string;
metadataHash?: string;
pageContent: string;
metadata: Metadata;
calculateHashes(): void;
toDocument(): DocumentInterface;
}
interface HashedDocumentArgs {
pageContent: string;
metadata: Metadata;
uid: string;
}
/**
* HashedDocument is a Document with hashes calculated.
* Hashes are calculated based on page content and metadata.
* It is used for indexing.
*/
export class _HashedDocument implements HashedDocumentInterface {
uid: string;
hash_?: string;
contentHash?: string;
metadataHash?: string;
pageContent: string;
metadata: Metadata;
constructor(fields: HashedDocumentArgs) {
this.uid = fields.uid;
this.pageContent = fields.pageContent;
this.metadata = fields.metadata;
}
calculateHashes(): void {
const forbiddenKeys = ["hash_", "content_hash", "metadata_hash"];
for (const key of forbiddenKeys) {
if (key in this.metadata) {
throw new Error(
`Metadata cannot contain key ${key} as it is reserved for internal use. Restricted keys: [${forbiddenKeys.join(
", "
)}]`
);
}
}
const contentHash = this._hashStringToUUID(this.pageContent);
try {
const metadataHash = this._hashNestedDictToUUID(this.metadata);
this.contentHash = contentHash;
this.metadataHash = metadataHash;
} catch (e) {
throw new Error(
`Failed to hash metadata: ${e}. Please use a dict that can be serialized using json.`
);
}
this.hash_ = this._hashStringToUUID(this.contentHash + this.metadataHash);
if (!this.uid) {
this.uid = this.hash_;
}
}
toDocument(): DocumentInterface {
return new Document({
pageContent: this.pageContent,
metadata: this.metadata,
});
}
static fromDocument(
document: DocumentInterface,
uid?: string
): _HashedDocument {
const doc = new this({
pageContent: document.pageContent,
metadata: document.metadata,
uid: uid || (document as DocumentInterface & { uid: string }).uid,
});
doc.calculateHashes();
return doc;
}
private _hashStringToUUID(inputString: string): string {
const hash_value = insecureHash(inputString);
return uuidv5(hash_value, UUIDV5_NAMESPACE);
}
private _hashNestedDictToUUID(data: Record<string, unknown>): string {
const serialized_data = JSON.stringify(data, Object.keys(data).sort());
const hash_value = insecureHash(serialized_data);
return uuidv5(hash_value, UUIDV5_NAMESPACE);
}
}
export type CleanupMode = "full" | "incremental";
export type IndexOptions = {
/**
* The number of documents to index in one batch.
*/
batchSize?: number;
/**
* The cleanup mode to use. Can be "full", "incremental" or undefined.
* - **Incremental**: Cleans up all documents that haven't been updated AND
* that are associated with source ids that were seen
* during indexing.
* Clean up is done continuously during indexing helping
* to minimize the probability of users seeing duplicated
* content.
* - **Full**: Delete all documents that haven to been returned by the loader.
* Clean up runs after all documents have been indexed.
* This means that users may see duplicated content during indexing.
* - **undefined**: Do not delete any documents.
*/
cleanup?: CleanupMode;
/**
* Optional key that helps identify the original source of the document.
* Must either be a string representing the key of the source in the metadata
* or a function that takes a document and returns a string representing the source.
* **Required when cleanup is incremental**.
*/
sourceIdKey?: StringOrDocFunc;
/**
* Batch size to use when cleaning up documents.
*/
cleanupBatchSize?: number;
/**
* Force update documents even if they are present in the
* record manager. Useful if you are re-indexing with updated embeddings.
*/
forceUpdate?: boolean;
};
export function _batch<T>(size: number, iterable: T[]): T[][] {
const batches: T[][] = [];
let currentBatch: T[] = [];
iterable.forEach((item) => {
currentBatch.push(item);
if (currentBatch.length >= size) {
batches.push(currentBatch);
currentBatch = [];
}
});
if (currentBatch.length > 0) {
batches.push(currentBatch);
}
return batches;
}
export function _deduplicateInOrder(
hashedDocuments: HashedDocumentInterface[]
): HashedDocumentInterface[] {
const seen = new Set<string>();
const deduplicated: HashedDocumentInterface[] = [];
for (const hashedDoc of hashedDocuments) {
if (!hashedDoc.hash_) {
throw new Error("Hashed document does not have a hash");
}
if (!seen.has(hashedDoc.hash_)) {
seen.add(hashedDoc.hash_);
deduplicated.push(hashedDoc);
}
}
return deduplicated;
}
export function _getSourceIdAssigner(
sourceIdKey: StringOrDocFunc | null
): (doc: DocumentInterface) => string | null {
if (sourceIdKey === null) {
return (_doc: DocumentInterface) => null;
} else if (typeof sourceIdKey === "string") {
return (doc: DocumentInterface) => doc.metadata[sourceIdKey];
} else if (typeof sourceIdKey === "function") {
return sourceIdKey;
} else {
throw new Error(
`sourceIdKey should be null, a string or a function, got ${typeof sourceIdKey}`
);
}
}
// eslint-disable-next-line @typescript-eslint/no-explicit-any
export const _isBaseDocumentLoader = (arg: any): arg is BaseDocumentLoader => {
if (
"load" in arg &&
typeof arg.load === "function" &&
"loadAndSplit" in arg &&
typeof arg.loadAndSplit === "function"
) {
return true;
}
return false;
};
interface IndexArgs {
docsSource: BaseDocumentLoader | DocumentInterface[];
recordManager: RecordManagerInterface;
vectorStore: VectorStore;
options?: IndexOptions;
}
/**
* Index data from the doc source into the vector store.
*
* Indexing functionality uses a manager to keep track of which documents
* are in the vector store.
*
* This allows us to keep track of which documents were updated, and which
* documents were deleted, which documents should be skipped.
*
* For the time being, documents are indexed using their hashes, and users
* are not able to specify the uid of the document.
*
* @param {IndexArgs} args
* @param {BaseDocumentLoader | DocumentInterface[]} args.docsSource The source of documents to index. Can be a DocumentLoader or a list of Documents.
* @param {RecordManagerInterface} args.recordManager The record manager to use for keeping track of indexed documents.
* @param {VectorStore} args.vectorStore The vector store to use for storing the documents.
* @param {IndexOptions | undefined} args.options Options for indexing.
* @returns {Promise<IndexingResult>}
*/
export async function index(args: IndexArgs): Promise<IndexingResult> {
const { docsSource, recordManager, vectorStore, options } = args;
const {
batchSize = 100,
cleanup,
sourceIdKey,
cleanupBatchSize = 1000,
forceUpdate = false,
} = options ?? {};
if (cleanup === "incremental" && !sourceIdKey) {
throw new Error(
"sourceIdKey is required when cleanup mode is incremental. Please provide through 'options.sourceIdKey'."
);
}
const docs = _isBaseDocumentLoader(docsSource)
? await docsSource.load()
: docsSource;
const sourceIdAssigner = _getSourceIdAssigner(sourceIdKey ?? null);
const indexStartDt = await recordManager.getTime();
let numAdded = 0;
let numDeleted = 0;
let numUpdated = 0;
let numSkipped = 0;
const batches = _batch<DocumentInterface>(batchSize ?? 100, docs);
for (const batch of batches) {
const hashedDocs = _deduplicateInOrder(
batch.map((doc) => _HashedDocument.fromDocument(doc))
);
const sourceIds = hashedDocs.map((doc) => sourceIdAssigner(doc));
if (cleanup === "incremental") {
hashedDocs.forEach((_hashedDoc, index) => {
const source = sourceIds[index];
if (source === null) {
throw new Error(
"sourceIdKey must be provided when cleanup is incremental"
);
}
});
}
const batchExists = await recordManager.exists(
hashedDocs.map((doc) => doc.uid)
);
const uids: string[] = [];
const docsToIndex: DocumentInterface[] = [];
const docsToUpdate: string[] = [];
const seenDocs = new Set<string>();
hashedDocs.forEach((hashedDoc, i) => {
const docExists = batchExists[i];
if (docExists) {
if (forceUpdate) {
seenDocs.add(hashedDoc.uid);
} else {
docsToUpdate.push(hashedDoc.uid);
return;
}
}
uids.push(hashedDoc.uid);
docsToIndex.push(hashedDoc.toDocument());
});
if (docsToUpdate.length > 0) {
await recordManager.update(docsToUpdate, { timeAtLeast: indexStartDt });
numSkipped += docsToUpdate.length;
}
if (docsToIndex.length > 0) {
await vectorStore.addDocuments(docsToIndex, { ids: uids });
numAdded += docsToIndex.length - seenDocs.size;
numUpdated += seenDocs.size;
}
await recordManager.update(
hashedDocs.map((doc) => doc.uid),
{ timeAtLeast: indexStartDt, groupIds: sourceIds }
);
if (cleanup === "incremental") {
sourceIds.forEach((sourceId) => {
if (!sourceId) throw new Error("Source id cannot be null");
});
const uidsToDelete = await recordManager.listKeys({
before: indexStartDt,
groupIds: sourceIds,
});
await vectorStore.delete({ ids: uidsToDelete });
await recordManager.deleteKeys(uidsToDelete);
numDeleted += uidsToDelete.length;
}
}
if (cleanup === "full") {
let uidsToDelete = await recordManager.listKeys({
before: indexStartDt,
limit: cleanupBatchSize,
});
while (uidsToDelete.length > 0) {
await vectorStore.delete({ ids: uidsToDelete });
await recordManager.deleteKeys(uidsToDelete);
numDeleted += uidsToDelete.length;
uidsToDelete = await recordManager.listKeys({
before: indexStartDt,
limit: cleanupBatchSize,
});
}
}
return {
numAdded,
numDeleted,
numUpdated,
numSkipped,
};
}