refactor: improve DX using image vector stores

run-llama · Dec 18, 2023 · 76f7256 · 76f7256
1 parent e381f95
commit 76f7256
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 20 deletions.
diff --git a/examples/multimodal/load.ts b/examples/multimodal/load.ts
@@ -2,7 +2,6 @@ import {
   ServiceContext,
   serviceContextFromDefaults,
   SimpleDirectoryReader,
-  SimpleVectorStore,
   storageContextFromDefaults,
   VectorStoreIndex,
 } from "llamaindex";
@@ -22,17 +21,9 @@ async function generateDatasource(serviceContext: ServiceContext) {
     const documents = await new SimpleDirectoryReader().loadData({
       directoryPath: path.join("multimodal", "data"),
     });
-    // set up vector store index with two vector stores, one for text, the other for images
-    const vectorStore = await SimpleVectorStore.fromPersistDir(
-      path.join("storage", "text"),
-    );
-    const imageVectorStore = await SimpleVectorStore.fromPersistDir(
-      path.join("storage", "images"),
-    );
     const storageContext = await storageContextFromDefaults({
       persistDir: "storage",
-      vectorStore,
-      imageVectorStore,
+      storeImages: true,
     });
     await VectorStoreIndex.fromDocuments(documents, {
       serviceContext,

diff --git a/examples/multimodal/retrieve.ts b/examples/multimodal/retrieve.ts
@@ -1,7 +1,6 @@
 import {
   ImageNode,
   serviceContextFromDefaults,
-  SimpleVectorStore,
   storageContextFromDefaults,
   TextNode,
   VectorStoreIndex,
@@ -14,16 +13,9 @@ export async function createRetriever() {
     chunkSize: 512,
     chunkOverlap: 20,
   });
-  const vectorStore = await SimpleVectorStore.fromPersistDir(
-    path.join("storage", "text"),
-  );
-  const imageVectorStore = await SimpleVectorStore.fromPersistDir(
-    path.join("storage", "images"),
-  );
   const storageContext = await storageContextFromDefaults({
     persistDir: "storage",
-    vectorStore,
-    imageVectorStore,
+    storeImages: true,
   });
   const index = await VectorStoreIndex.init({
     nodes: [],

diff --git a/packages/core/src/storage/StorageContext.ts b/packages/core/src/storage/StorageContext.ts
@@ -1,5 +1,10 @@
+import * as path from "path";
 import { GenericFileSystem } from "./FileSystem";
-import { DEFAULT_FS, DEFAULT_NAMESPACE } from "./constants";
+import {
+  DEFAULT_FS,
+  DEFAULT_IMAGE_VECTOR_NAMESPACE,
+  DEFAULT_NAMESPACE,
+} from "./constants";
 import { SimpleDocumentStore } from "./docStore/SimpleDocumentStore";
 import { BaseDocumentStore } from "./docStore/types";
 import { SimpleIndexStore } from "./indexStore/SimpleIndexStore";
@@ -19,6 +24,7 @@ type BuilderParams = {
   indexStore?: BaseIndexStore;
   vectorStore?: VectorStore;
   imageVectorStore?: VectorStore;
+  storeImages?: boolean;
   persistDir?: string;
   fs?: GenericFileSystem;
 };
@@ -28,13 +34,15 @@ export async function storageContextFromDefaults({
   indexStore,
   vectorStore,
   imageVectorStore,
+  storeImages,
   persistDir,
   fs,
 }: 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 || DEFAULT_FS;
     docStore =
@@ -48,6 +56,12 @@ export async function storageContextFromDefaults({
       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 {

diff --git a/packages/core/src/storage/constants.ts b/packages/core/src/storage/constants.ts
@@ -5,4 +5,5 @@ export const DEFAULT_DOC_STORE_PERSIST_FILENAME = "doc_store.json";
 export const DEFAULT_VECTOR_STORE_PERSIST_FILENAME = "vector_store.json";
 export const DEFAULT_GRAPH_STORE_PERSIST_FILENAME = "graph_store.json";
 export const DEFAULT_NAMESPACE = "docstore";
+export const DEFAULT_IMAGE_VECTOR_NAMESPACE = "images";
 export { DEFAULT_FS } from "./FileSystem";