From b8c4bd70aec5826f9df1f0c8ca086f631deeec4e Mon Sep 17 00:00:00 2001
From: Marcus Schiesser <mail@marcusschiesser.de>
Date: Thu, 14 Dec 2023 10:45:07 +0700
Subject: [PATCH] refactor: improve DX using image vector stores

---
 examples/multimodal/load.ts                 | 11 +----------
 examples/multimodal/retrieve.ts             | 10 +---------
 packages/core/src/storage/StorageContext.ts | 16 +++++++++++++++-
 packages/core/src/storage/constants.ts      |  1 +
 4 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/examples/multimodal/load.ts b/examples/multimodal/load.ts
index 5c0a097c3..48556686e 100644
--- 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
index 6b3c83b37..9366935a4 100644
--- 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
index 8e557ea95..9163677c6 100644
--- 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,6 +34,7 @@ export async function storageContextFromDefaults({
   indexStore,
   vectorStore,
   imageVectorStore,
+  storeImages,
   persistDir,
   fs,
 }: BuilderParams): Promise<StorageContext> {
@@ -35,6 +42,7 @@ export async function storageContextFromDefaults({
     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
index 15e87613c..c4d965d2b 100644
--- 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";
-- 
GitLab