-
Notifications
You must be signed in to change notification settings - Fork 295
/
StorageContext.ts
73 lines (69 loc) · 2.32 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
73
import { path } from "@llamaindex/env";
import { ModalityType, ObjectType } from "../Node.js";
import { ClipEmbedding } from "../embeddings/ClipEmbedding.js";
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, VectorStoreByType } from "./vectorStore/types.js";
export interface StorageContext {
docStore: BaseDocumentStore;
indexStore: BaseIndexStore;
vectorStores: VectorStoreByType;
}
type BuilderParams = {
docStore: BaseDocumentStore;
indexStore: BaseIndexStore;
vectorStore: VectorStore;
vectorStores: VectorStoreByType;
storeImages: boolean;
persistDir: string;
};
export async function storageContextFromDefaults({
docStore,
indexStore,
vectorStore,
vectorStores,
storeImages,
persistDir,
}: Partial<BuilderParams>): Promise<StorageContext> {
vectorStores = vectorStores ?? {};
if (!persistDir) {
docStore = docStore ?? new SimpleDocumentStore();
indexStore = indexStore ?? new SimpleIndexStore();
if (!(ModalityType.TEXT in vectorStores)) {
vectorStores[ModalityType.TEXT] = vectorStore ?? new SimpleVectorStore();
}
if (storeImages && !(ModalityType.IMAGE in vectorStores)) {
vectorStores[ModalityType.IMAGE] = new SimpleVectorStore({
embedModel: new ClipEmbedding(),
});
}
} else {
docStore =
docStore ||
(await SimpleDocumentStore.fromPersistDir(persistDir, DEFAULT_NAMESPACE));
indexStore =
indexStore || (await SimpleIndexStore.fromPersistDir(persistDir));
if (!(ObjectType.TEXT in vectorStores)) {
vectorStores[ModalityType.TEXT] =
vectorStore ?? (await SimpleVectorStore.fromPersistDir(persistDir));
}
if (storeImages && !(ObjectType.IMAGE in vectorStores)) {
vectorStores[ModalityType.IMAGE] = await SimpleVectorStore.fromPersistDir(
path.join(persistDir, DEFAULT_IMAGE_VECTOR_NAMESPACE),
new ClipEmbedding(),
);
}
}
return {
docStore,
indexStore,
vectorStores,
};
}