vanna 2.0.1__py3-none-any.whl → 2.0.2__py3-none-any.whl

vanna/integrations/chromadb/agent_memory.py

@@ -16,6 +16,15 @@ try:
     from chromadb.config import Settings
     from chromadb.utils import embedding_functions
 
+    try:
+        from chromadb.errors import NotFoundError
+    except ImportError:
+        # Fallback for older ChromaDB versions that don't have chromadb.errors
+        class NotFoundError(Exception):
+            """Fallback NotFoundError for older ChromaDB versions."""
+
+            pass
+
     CHROMADB_AVAILABLE = True
 except ImportError:
     CHROMADB_AVAILABLE = False
@@ -31,7 +40,64 @@ from vanna.core.tool import ToolContext
 
 
 class ChromaAgentMemory(AgentMemory):
-    """ChromaDB-based implementation of AgentMemory."""
+    """ChromaDB-based implementation of AgentMemory.
+
+    This implementation uses ChromaDB's PersistentClient to store agent memories
+    on disk, ensuring they persist across application restarts.
+
+    Key Features:
+    - Persistent storage: All memories are automatically saved to disk
+    - Efficient retrieval: Existing collections are loaded without re-initializing
+      embedding functions, avoiding unnecessary model downloads
+    - Flexible embedding: Supports custom embedding functions or uses ChromaDB's
+      default embedding function
+
+    Args:
+        persist_directory: Directory where ChromaDB will store its data.
+                          Defaults to "./chroma_memory". Use an absolute path
+                          for production deployments to ensure consistent location
+                          across restarts.
+        collection_name: Name of the ChromaDB collection to use. Multiple agents
+                        can share the same persist_directory with different
+                        collection names.
+        embedding_function: Optional custom embedding function. If not provided,
+                           ChromaDB's DefaultEmbeddingFunction is used (requires
+                           internet connection on first use to download the model).
+                           Once a collection is created, subsequent application
+                           restarts will retrieve the existing collection without
+                           re-downloading the model.
+
+    Example:
+        >>> from vanna.integrations.chromadb import ChromaAgentMemory
+        >>> # Basic usage with defaults
+        >>> memory = ChromaAgentMemory(
+        ...     persist_directory="/app/data/chroma",
+        ...     collection_name="my_agent_memory"
+        ... )
+        >>>
+        >>> # With custom embedding function (e.g., for offline use)
+        >>> from chromadb.utils import embedding_functions
+        >>> ef = embedding_functions.SentenceTransformerEmbeddingFunction()
+        >>> memory = ChromaAgentMemory(
+        ...     persist_directory="/app/data/chroma",
+        ...     embedding_function=ef
+        ... )
+
+    Note:
+        The default embedding function downloads an ONNX model (~80MB) on first use.
+        For air-gapped or offline environments, pre-download the model or provide
+        a custom embedding function.
+
+    Limitation:
+        This class does not validate that an existing Chroma collection was created
+        with the same embedding function as the one configured for the current
+        ``ChromaAgentMemory`` instance. If you reuse a collection (same
+        ``persist_directory`` and ``collection_name``) with a different embedding
+        function than was originally used, queries may fail or produce incorrect
+        similarity results. It is your responsibility to ensure that a given
+        collection is always accessed with a consistent embedding function, or to
+        implement your own validation around collection creation and reuse.
+    """
 
     def __init__(
         self,
@@ -76,12 +142,14 @@ class ChromaAgentMemory(AgentMemory):
         """Get or create ChromaDB collection."""
         if self._collection is None:
             client = self._get_client()
-            embedding_func = self._get_embedding_function()
             try:
-                self._collection = client.get_collection(
-                    name=self.collection_name, embedding_function=embedding_func
-                )
-            except Exception:
+                # Try to get existing collection first
+                # Don't pass embedding_function here to avoid re-instantiating/downloading it
+                # For existing collections, ChromaDB uses the stored embedding function configuration
+                self._collection = client.get_collection(name=self.collection_name)
+            except NotFoundError:
+                # Collection doesn't exist, create it with embedding function
+                embedding_func = self._get_embedding_function()
                 self._collection = client.create_collection(
                     name=self.collection_name,
                     embedding_function=embedding_func,

{vanna-2.0.1.dist-info → vanna-2.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vanna
-Version: 2.0.1
+Version: 2.0.2
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9

{vanna-2.0.1.dist-info → vanna-2.0.2.dist-info}/RECORD

@@ -128,7 +128,7 @@ vanna/integrations/azuresearch/agent_memory.py,sha256=zzGQmPOgrkLy5PA63joDa5Eprr
 vanna/integrations/bigquery/__init__.py,sha256=Sw-krShPv3nOLQIfVQSflCQWXqIQ-gF_eCuE6GHatNY,108
 vanna/integrations/bigquery/sql_runner.py,sha256=65CcGdX6kxOS5wUGSQ558w4_0oH412tV0f5jVwDczx0,2691
 vanna/integrations/chromadb/__init__.py,sha256=sPumYrf9hMVgEHh6rzXdWu9OrVkOucA5_HkQ4LXfIUY,3642
-vanna/integrations/chromadb/agent_memory.py,sha256=3s06jDc4v1m6MHSMLM3PKza9Em0Pm2hE6NMbGqD8HmU,14833
+vanna/integrations/chromadb/agent_memory.py,sha256=nccMdv2VDrFkefw-hfKRLFUC_TgruApmKzR4T6lvBKo,18347
 vanna/integrations/clickhouse/__init__.py,sha256=Ke7P4HrHtgJZTfCnKv2P8flAS6IY-luZLtJCZmDlCig,114
 vanna/integrations/clickhouse/sql_runner.py,sha256=195z-Uij69AwyPOW3nO6ipjgGfA0oJ1TgxLcsGa7VtE,2350
 vanna/integrations/duckdb/__init__.py,sha256=kQiTv0Z5ozcIpL8FaZpZeYfxgbE6eteqIwjlbaqV4Ig,102
@@ -282,8 +282,8 @@ vanna/tools/run_sql.py,sha256=RGptT_cRfD6FP-3MnTS_pgYahWJbDSb3gv3yZ-5sR3g,6951
 vanna/tools/visualize_data.py,sha256=qgPMTxEAmbtJjb2j-qtAG1ZwtzvBvd8gdkzGMki_YGE,7551
 vanna/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vanna/web_components/__init__.py,sha256=QLzEJU3ersllv4pgZpaKdws3S3u5vDyTyUivawMc4A0,1040
-vanna-2.0.1.dist-info/entry_points.txt,sha256=lWU4rLAVNmrWnY6E66z_2TQW3QHU2Kzhwu2qvC6EBrs,62
-vanna-2.0.1.dist-info/licenses/LICENSE,sha256=VYiPMMDqj9BcxUkAYqrAzJpn5tCFXCrnglfRqS5l9Rk,1065
-vanna-2.0.1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-vanna-2.0.1.dist-info/METADATA,sha256=6eTDvS_ohU3WNtJwlBvGg1doY2hqDZmt7xnPjuV5Hrs,16519
-vanna-2.0.1.dist-info/RECORD,,
+vanna-2.0.2.dist-info/entry_points.txt,sha256=lWU4rLAVNmrWnY6E66z_2TQW3QHU2Kzhwu2qvC6EBrs,62
+vanna-2.0.2.dist-info/licenses/LICENSE,sha256=VYiPMMDqj9BcxUkAYqrAzJpn5tCFXCrnglfRqS5l9Rk,1065
+vanna-2.0.2.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+vanna-2.0.2.dist-info/METADATA,sha256=Go4qGzf-AZrnvDZd0pPUJkfH_0U-v2qy08yRv1j5Uf8,16519
+vanna-2.0.2.dist-info/RECORD,,