-
Notifications
You must be signed in to change notification settings - Fork 318
/
StorageContext.ts
72 lines (68 loc) · 2.07 KB
/
StorageContext.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
import type { GenericFileSystem } from "@llamaindex/env";
import { defaultFS, path } from "@llamaindex/env";
import {
DEFAULT_IMAGE_VECTOR_NAMESPACE,
DEFAULT_NAMESPACE,
} from "./constants.js";
import { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
import type { BaseDocumentStore } from "./docStore/types.js";
import { SimpleIndexStore } from "./indexStore/SimpleIndexStore.js";
import type { BaseIndexStore } from "./indexStore/types.js";
import { SimpleVectorStore } from "./vectorStore/SimpleVectorStore.js";
import type { VectorStore } from "./vectorStore/types.js";
export interface StorageContext {
docStore: BaseDocumentStore;
indexStore: BaseIndexStore;
vectorStore: VectorStore;
imageVectorStore?: VectorStore;
}
export type BuilderParams = {
docStore: BaseDocumentStore;
indexStore: BaseIndexStore;
vectorStore: VectorStore;
imageVectorStore: VectorStore;
storeImages: boolean;
persistDir: string;
fs: GenericFileSystem;
};
export async function storageContextFromDefaults({
docStore,
indexStore,
vectorStore,
imageVectorStore,
storeImages,
persistDir,
fs,
}: Partial<BuilderParams>): Promise<StorageContext> {
if (!persistDir) {
docStore = docStore || new SimpleDocumentStore();
indexStore = indexStore || new SimpleIndexStore();
vectorStore = vectorStore || new SimpleVectorStore();
imageVectorStore = storeImages ? new SimpleVectorStore() : imageVectorStore;
} else {
fs = fs || defaultFS;
docStore =
docStore ||
(await SimpleDocumentStore.fromPersistDir(
persistDir,
DEFAULT_NAMESPACE,
fs,
));
indexStore =
indexStore || (await SimpleIndexStore.fromPersistDir(persistDir, fs));
vectorStore =
vectorStore || (await SimpleVectorStore.fromPersistDir(persistDir, fs));
imageVectorStore = storeImages
? await SimpleVectorStore.fromPersistDir(
path.join(persistDir, DEFAULT_IMAGE_VECTOR_NAMESPACE),
fs,
)
: imageVectorStore;
}
return {
docStore,
indexStore,
vectorStore,
imageVectorStore,
};
}