[AUR-362] Add In-memory vector store (#22)

* [AUR-362] Add In-memory vector store

* [AUR-362] fix delete fun input format

* [AUR-362] revise persist and from persist path to save and load

* [AUR-362] revise simple.py to in_memory.py

knowledgehub/vectorstores/__init__.py

@@ -1,4 +1,5 @@
 from .base import BaseVectorStore
 from .chroma import ChromaVectorStore
+from .in_memory import InMemoryVectorStore
 
-__all__ = ["BaseVectorStore", "ChromaVectorStore"]
+__all__ = ["BaseVectorStore", "ChromaVectorStore", "InMemoryVectorStore"]

knowledgehub/vectorstores/base.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, List, Optional, Tuple, Type, Union
 
+from llama_index.schema import NodeRelationship, RelatedNodeInfo
 from llama_index.vector_stores.types import BasePydanticVectorStore
 from llama_index.vector_stores.types import VectorStore as LIVectorStore
 from llama_index.vector_stores.types import VectorStoreQuery
@@ -118,6 +119,9 @@ class LlamaIndexVectorStore(BaseVectorStore):
         if ids is not None:
             for node, id in zip(nodes, ids):
                 node.id_ = id
+                node.relationships = {
+                    NodeRelationship.SOURCE: RelatedNodeInfo(node_id=id)
+                }
 
         return self._client.add(nodes=nodes)  # type: ignore
 

knowledgehub/vectorstores/in_memory.py

@@ -0,0 +1,57 @@
+"""Simple vector store index."""
+
+from typing import Any, Optional, Type
+
+import fsspec
+from llama_index.vector_stores import SimpleVectorStore as LISimpleVectorStore
+from llama_index.vector_stores.simple import SimpleVectorStoreData
+
+from kotaemon.vectorstores.base import LlamaIndexVectorStore
+
+
+class InMemoryVectorStore(LlamaIndexVectorStore):
+    _li_class: Type[LISimpleVectorStore] = LISimpleVectorStore
+    store_text: bool = False
+
+    def __init__(
+        self,
+        data: Optional[SimpleVectorStoreData] = None,
+        fs: Optional[fsspec.AbstractFileSystem] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize params."""
+        self._data = data or SimpleVectorStoreData()
+        self._fs = fs or fsspec.filesystem("file")
+
+        super().__init__(
+            data=data,
+            fs=fs,
+            **kwargs,
+        )
+
+    def save(
+        self,
+        save_path: str,
+        fs: Optional[fsspec.AbstractFileSystem] = None,
+        **kwargs,
+    ):
+
+        """save a simpleVectorStore to a dictionary.
+
+        Args:
+            save_path: Path of saving vector to disk.
+            fs: An abstract super-class for pythonic file-systems
+        """
+        self._client.persist(persist_path=save_path, fs=fs)
+
+    def load(
+        self, load_path: str, fs: Optional[fsspec.AbstractFileSystem] = None
+    ) -> "InMemoryVectorStore":
+
+        """Create a SimpleKVStore from a load directory.
+
+        Args:
+            load_path: Path of loading vector.
+            fs: An abstract super-class for pythonic file-systems
+        """
+        return self._client.from_persist_path(persist_path=load_path, fs=fs)