swarmauri_vectorstore_persistentchromadb 0.6.0.dev154__tar.gz

swarmauri_vectorstore_persistentchromadb-0.6.0.dev154/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.3
+Name: swarmauri_vectorstore_persistentchromadb
+Version: 0.6.0.dev154
+Summary: A Persistent ChromaDB based Vector Store
+License: Apache-2.0
+Author: Jacob Stewart
+Author-email: jacob@swarmauri.com
+Requires-Python: >=3.10,<3.13
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: chromadb (>=0.5.17,<0.6.0)
+Requires-Dist: swarmauri_base (>=0.6.0.dev154,<0.7.0)
+Requires-Dist: swarmauri_core (>=0.6.0.dev154,<0.7.0)
+Requires-Dist: swarmauri_embedding_doc2vec (>=0.6.0.dev154,<0.7.0)
+Project-URL: Repository, http://github.com/swarmauri/swarmauri-sdk
+Description-Content-Type: text/markdown
+
+# Swarmauri Example Community Package

swarmauri_vectorstore_persistentchromadb-0.6.0.dev154/README.md

@@ -0,0 +1 @@
+# Swarmauri Example Community Package

swarmauri_vectorstore_persistentchromadb-0.6.0.dev154/pyproject.toml

@@ -0,0 +1,58 @@
+[tool.poetry]
+name = "swarmauri_vectorstore_persistentchromadb"
+version = "0.6.0.dev154"
+description = "A Persistent ChromaDB based Vector Store"
+authors = ["Jacob Stewart <jacob@swarmauri.com>"]
+license = "Apache-2.0"
+readme = "README.md"
+repository = "http://github.com/swarmauri/swarmauri-sdk"
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12"
+]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<3.13"
+
+# Swarmauri
+swarmauri_core = {version = "^0.6.0.dev154"}
+swarmauri_base = {version = "^0.6.0.dev154"}
+swarmauri_embedding_doc2vec = {version = "^0.6.0.dev154"}
+
+
+# Dependencies
+chromadb = "^0.5.17"
+
+[tool.poetry.group.dev.dependencies]
+flake8 = "^7.0"
+pytest = "^8.0"
+pytest-asyncio = ">=0.24.0"
+pytest-xdist = "^3.6.1"
+pytest-json-report = "^1.5.0"
+python-dotenv = "*"
+requests = "^2.32.3"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+norecursedirs = ["combined", "scripts"]
+
+markers = [
+    "test: standard test",
+    "unit: Unit tests",
+    "integration: Integration tests",
+    "acceptance: Acceptance tests",
+    "experimental: Experimental tests"
+]
+log_cli = true
+log_cli_level = "INFO"
+log_cli_format = "%(asctime)s [%(levelname)s] %(message)s"
+log_cli_date_format = "%Y-%m-%d %H:%M:%S"
+asyncio_default_fixture_loop_scope = "function"
+
+[tool.poetry.plugins."swarmauri.vector_stores"]
+PersistentChromaDBVectorStore = "swarmauri_vectorstore_persistentchromadb.PersistentChromaDBVectorStore:PersistentChromaDBVectorStore"

swarmauri_vectorstore_persistentchromadb-0.6.0.dev154/swarmauri_vectorstore_persistentchromadb/PersistentChromaDBVectorStore.py

@@ -0,0 +1,188 @@
+import logging
+import chromadb
+
+from typing import List, Union, Literal
+
+from swarmauri_standard.documents.concrete.Document import Document
+from swarmauri_embedding_doc2vec.Doc2VecEmbedding import Doc2VecEmbedding
+from swarmauri.distances.concrete.CosineDistance import CosineDistance
+
+from swarmauri.vector_stores.base.VectorStoreBase import VectorStoreBase
+from swarmauri.vector_stores.base.VectorStoreRetrieveMixin import (
+    VectorStoreRetrieveMixin,
+)
+from swarmauri.vector_stores.base.VectorStoreSaveLoadMixin import (
+    VectorStoreSaveLoadMixin,
+)
+from swarmauri.vector_stores.base.VectorStorePersistentMixin import (
+    VectorStorePersistentMixin,
+)
+from swarmauri_core.ComponentBase import ComponentBase
+
+
+@ComponentBase.register_type(VectorStoreBase, "PersistentChromaDBVectorStore")
+class PersistentChromaDBVectorStore(
+    VectorStoreSaveLoadMixin,
+    VectorStoreRetrieveMixin,
+    VectorStorePersistentMixin,
+    VectorStoreBase,
+):
+    type: Literal["PersistentChromaDBVectorStore"] = "PersistentChromaDBVectorStore"
+
+    def __init__(self, **kwargs):
+        """
+        Initialize the PersistentChromaDBVectorStore.
+
+        Args:
+            Args:
+            **kwargs: keyword arguments.
+        """
+        super().__init__(**kwargs)
+
+        self._embedder = Doc2VecEmbedding(vector_size=self.vector_size)
+        self._distance = CosineDistance()
+
+    def connect(self) -> None:
+        """
+        Establish a connection to ChromaDB and get or create the collection.
+        """
+        # settings = Settings(
+        #     chroma_api_impl="chromadb.api.fastapi.FastAPI",  # Use FastAPI if LocalAPI is not supported
+        #     chroma_server_host="localhost",  # Server host
+        #     chroma_server_http_port=8000,  # Server port
+        # )
+        #
+        # self.client = chromadb.Client(
+        #     settings=settings,
+        # )
+        self.client = chromadb.Client()
+
+        self.collection = self.client.get_or_create_collection(
+            name=self.collection_name
+        )
+        logging.info(
+            f"Connected to ChromaDB at {self.path}, collection: {self.collection_name}"
+        )
+
+    def disconnect(self) -> None:
+        """
+        Close the connection to ChromaDB.
+        """
+        if self.client:
+            # Perform any necessary cleanup here
+            self.client = None
+            self.collection = None
+
+    def add_document(self, document: Document) -> None:
+        embedding = None
+        if not document.embedding:
+            self._embedder.fit([document.content])  # Fit only once
+            embedding = (
+                self._embedder.transform([document.content])[0].to_numpy().tolist()
+            )
+        else:
+            embedding = document.embedding
+
+        self.collection.add(
+            ids=[document.id],
+            documents=[document.content],
+            embeddings=[embedding],
+            metadatas=[document.metadata],
+        )
+
+    def add_documents(self, documents: List[Document]) -> None:
+        ids = [doc.id for doc in documents]
+        texts = [doc.content for doc in documents]
+
+        for doc in documents:
+            self._embedder.fit([doc.content])
+
+        embeddings = [
+            self._embedder.infer_vector(doc.content).value for doc in documents
+        ]
+        metadatas = [doc.metadata for doc in documents]
+        if metadatas[0]:
+            self.collection.add(
+                ids=ids, documents=texts, embeddings=embeddings, metadatas=metadatas
+            )
+        else:
+            self.collection.add(ids=ids, documents=texts, embeddings=embeddings)
+
+    def get_document(self, doc_id: str) -> Union[Document, None]:
+        results = self.collection.get(ids=[doc_id])
+        if not results["metadatas"][0]:
+            results["metadatas"][0] = {}
+        if results["ids"]:
+            document = Document(
+                id=results["ids"][0],
+                content=results["documents"][0],
+                metadata=results["metadatas"][0],
+            )
+            return document
+        return None
+
+    def get_all_documents(self) -> List[Document]:
+        results = self.collection.get()
+        documents = [
+            Document(
+                id=results["ids"][idx],
+                content=results["documents"][idx],
+                metadata=results["metadatas"][idx],
+            )
+            for idx in range(len(results["ids"]))
+        ]
+        return documents
+
+    def delete_document(self, doc_id: str) -> None:
+        self.collection.delete(ids=[doc_id])
+
+    def update_document(self, doc_id: str, updated_document: Document) -> None:
+        document_vector = None
+        # Precompute the embedding outside the update process
+        if not updated_document.embedding:
+            # Transform without refitting to avoid vocabulary issues
+            document_vector = self._embedder.transform([updated_document.content])[0]
+        else:
+            document_vector = updated_document.embedding
+
+        document_vector = document_vector.to_numpy().tolist()
+
+        updated_document.embedding = document_vector
+
+        self.delete_document(doc_id)
+        self.add_document(updated_document)
+
+    def clear_documents(self) -> None:
+        documents = self.get_all_documents()
+        doc_ids = [doc.id for doc in documents]
+        self.collection.delete(ids=doc_ids)
+
+    def document_count(self) -> int:
+        return len(self.get_all_documents())
+
+    def retrieve(self, query: str, top_k: int = 5) -> List[Document]:
+        query_embedding = self._embedder.infer_vector(query).value
+        # print(query_embedding)
+
+        results = self.collection.query(
+            query_embeddings=query_embedding, n_results=top_k
+        )
+
+        return [
+            Document(
+                id=results["ids"][0][idx],
+                content=results["documents"][0][idx],
+                metadata=(
+                    results["metadatas"][0][idx] if results["metadatas"][0][idx] else {}
+                ),
+            )
+            for idx in range(len(results["ids"][0]))
+        ]
+
+    # Override the model_dump_json method
+    def model_dump_json(self, *args, **kwargs) -> str:
+        # Call the disconnect method before serialization
+        self.disconnect()
+
+        # Now proceed with the usual JSON serialization
+        return super().model_dump_json(*args, **kwargs)

swarmauri_vectorstore_persistentchromadb-0.6.0.dev154/swarmauri_vectorstore_persistentchromadb/__init__.py

@@ -0,0 +1,14 @@
+from .PersistentChromaDBVectorStore import PersistentChromaDBVectorStore
+
+__version__ = "0.6.0.dev26"
+__long_desc__ = """
+
+# Swarmauri PersistentChromaDB Plugin
+
+This repository includes an PersistentChromaDB of a Swarmauri Plugin.
+
+Visit us at: https://swarmauri.com
+Follow us at: https://github.com/swarmauri
+Star us at: https://github.com/swarmauri/swarmauri-sdk
+
+"""