vanna 2.0.1__tar.gz → 2.0.2__tar.gz

{vanna-2.0.1 → vanna-2.0.2}/PKG-INFO

@@ -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 → vanna-2.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "2.0.1"
+version = "2.0.2"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -219,4 +219,4 @@ indent-style = "space"
 skip-magic-trailing-comma = false
 
 # Automatically detect line endings
-line-ending = "auto"
+line-ending = "auto"

{vanna-2.0.1 → vanna-2.0.2}/src/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,