Combine docstores and vectorstores within a storages component (#72)

knowledgehub/storages/vectorstores/chroma.py

@@ -0,0 +1,77 @@
+from typing import Any, Dict, List, Optional, Type, cast
+
+from llama_index.vector_stores.chroma import ChromaVectorStore as LIChromaVectorStore
+
+from .base import LlamaIndexVectorStore
+
+
+class ChromaVectorStore(LlamaIndexVectorStore):
+    _li_class: Type[LIChromaVectorStore] = LIChromaVectorStore
+
+    def __init__(
+        self,
+        path: str = "./chroma",
+        collection_name: str = "default",
+        host: str = "localhost",
+        port: str = "8000",
+        ssl: bool = False,
+        headers: Optional[Dict[str, str]] = None,
+        collection_kwargs: Optional[dict] = None,
+        stores_text: bool = True,
+        flat_metadata: bool = True,
+        **kwargs: Any,
+    ):
+        try:
+            import chromadb
+        except ImportError:
+            raise ImportError(
+                "ChromaVectorStore requires chromadb. "
+                "Please install chromadb first `pip install chromadb`"
+            )
+
+        client = chromadb.PersistentClient(path=path)
+        collection = client.get_or_create_collection(collection_name)
+
+        # pass through for nice IDE support
+        super().__init__(
+            chroma_collection=collection,
+            host=host,
+            port=port,
+            ssl=ssl,
+            headers=headers or {},
+            collection_kwargs=collection_kwargs or {},
+            stores_text=stores_text,
+            flat_metadata=flat_metadata,
+            **kwargs,
+        )
+        self._client = cast(LIChromaVectorStore, self._client)
+
+    def delete(self, ids: List[str], **kwargs):
+        """Delete vector embeddings from vector stores
+
+        Args:
+            ids: List of ids of the embeddings to be deleted
+            kwargs: meant for vectorstore-specific parameters
+        """
+        self._client._collection.delete(ids=ids)
+
+    def delete_collection(self, collection_name: Optional[str] = None):
+        """Delete entire collection under specified name from vector stores
+
+        Args:
+            collection_name: Name of the collection to delete
+        """
+        # a rather ugly chain call but it do the job of finding
+        # original chromadb client and call delete_collection() method
+        if collection_name is None:
+            collection_name = self._client.client.name
+        self._client.client._client.delete_collection(collection_name)
+
+    def count(self) -> int:
+        return self._collection.count()
+
+    def save(self, *args, **kwargs):
+        pass
+
+    def load(self, *args, **kwargs):
+        pass