agno 2.0.0rc2__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/vectordb/chroma/chromadb.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 from hashlib import md5
 from typing import Any, Dict, List, Mapping, Optional, Union, cast
 
@@ -12,10 +13,11 @@ try:
 except ImportError:
     raise ImportError("The `chromadb` package is not installed. Please install it via `pip install chromadb`.")
 
+from agno.filters import FilterExpr
 from agno.knowledge.document import Document
 from agno.knowledge.embedder import Embedder
 from agno.knowledge.reranker.base import Reranker
-from agno.utils.log import log_debug, log_error, log_info, logger
+from agno.utils.log import log_debug, log_error, log_info, log_warning, logger
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 
@@ -24,6 +26,9 @@ class ChromaDb(VectorDb):
     def __init__(
         self,
         collection: str,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        id: Optional[str] = None,
         embedder: Optional[Embedder] = None,
         distance: Distance = Distance.cosine,
         path: str = "tmp/chromadb",
@@ -31,9 +36,22 @@ class ChromaDb(VectorDb):
         reranker: Optional[Reranker] = None,
         **kwargs,
     ):
+        # Validate required parameters
+        if not collection:
+            raise ValueError("Collection name must be provided.")
+
+        # Dynamic ID generation based on unique identifiers
+        if id is None:
+            from agno.utils.string import generate_id
+
+            seed = f"{path}#{collection}"
+            id = generate_id(seed)
+
+        # Initialize base class with name, description, and generated ID
+        super().__init__(id=id, name=name, description=description)
+
         # Collection attributes
         self.collection_name: str = collection
-
         # Embedder for embedding the document contents
         if embedder is None:
             from agno.knowledge.embedder.openai import OpenAIEmbedder
@@ -60,6 +78,44 @@ class ChromaDb(VectorDb):
         # Chroma client kwargs
         self.kwargs = kwargs
 
+    def _flatten_metadata(self, metadata: Dict[str, Any]) -> Dict[str, Union[str, int, float, bool]]:
+        """
+        Flatten nested metadata to ChromaDB-compatible format.
+
+        Args:
+            metadata: Dictionary that may contain nested structures
+
+        Returns:
+            Flattened dictionary with only primitive values
+        """
+        flattened: Dict[str, Any] = {}
+
+        def _flatten_recursive(obj: Any, prefix: str = "") -> None:
+            if isinstance(obj, dict):
+                if len(obj) == 0:
+                    # Handle empty dictionaries by converting to JSON string
+                    flattened[prefix] = json.dumps(obj)
+                else:
+                    for key, value in obj.items():
+                        new_key = f"{prefix}.{key}" if prefix else key
+                        _flatten_recursive(value, new_key)
+            elif isinstance(obj, (list, tuple)):
+                # Convert lists/tuples to JSON strings
+                flattened[prefix] = json.dumps(obj)
+            elif isinstance(obj, (str, int, float, bool)) or obj is None:
+                if obj is not None:  # ChromaDB doesn't accept None values
+                    flattened[prefix] = obj
+            else:
+                # Convert other complex types to JSON strings
+                try:
+                    flattened[prefix] = json.dumps(obj)
+                except (TypeError, ValueError):
+                    # If it can't be serialized, convert to string
+                    flattened[prefix] = str(obj)
+
+        _flatten_recursive(metadata)
+        return flattened
+
     @property
     def client(self) -> ClientAPI:
         if self._client is None:
@@ -147,11 +203,14 @@ class ChromaDb(VectorDb):
 
             metadata["content_hash"] = content_hash
 
+            # Flatten metadata for ChromaDB compatibility
+            flattened_metadata = self._flatten_metadata(metadata)
+
             docs_embeddings.append(document.embedding)
             docs.append(cleaned_content)
             ids.append(doc_id)
-            docs_metadata.append(metadata)
-            log_debug(f"Prepared document: {document.id} | {document.name} | {metadata}")
+            docs_metadata.append(flattened_metadata)
+            log_debug(f"Prepared document: {document.id} | {document.name} | {flattened_metadata}")
 
         if self._collection is None:
             logger.warning("Collection does not exist")
@@ -173,11 +232,47 @@ class ChromaDb(VectorDb):
         if not self._collection:
             self._collection = self.client.get_collection(name=self.collection_name)
 
-        try:
-            embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
-            await asyncio.gather(*embed_tasks, return_exceptions=True)
-        except Exception as e:
-            log_error(f"Error processing document: {e}")
+        if self.embedder.enable_batch and hasattr(self.embedder, "async_get_embeddings_batch_and_usage"):
+            # Use batch embedding when enabled and supported
+            try:
+                # Extract content from all documents
+                doc_contents = [doc.content for doc in documents]
+
+                # Get batch embeddings and usage
+                embeddings, usages = await self.embedder.async_get_embeddings_batch_and_usage(doc_contents)
+
+                # Process documents with pre-computed embeddings
+                for j, doc in enumerate(documents):
+                    try:
+                        if j < len(embeddings):
+                            doc.embedding = embeddings[j]
+                            doc.usage = usages[j] if j < len(usages) else None
+                    except Exception as e:
+                        logger.error(f"Error assigning batch embedding to document '{doc.name}': {e}")
+
+            except Exception as e:
+                # Check if this is a rate limit error - don't fall back as it would make things worse
+                error_str = str(e).lower()
+                is_rate_limit = any(
+                    phrase in error_str
+                    for phrase in ["rate limit", "too many requests", "429", "trial key", "api calls / minute"]
+                )
+
+                if is_rate_limit:
+                    logger.error(f"Rate limit detected during batch embedding. {e}")
+                    raise e
+                else:
+                    logger.warning(f"Async batch embedding failed, falling back to individual embeddings: {e}")
+                    # Fall back to individual embedding
+                    embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in documents]
+                    await asyncio.gather(*embed_tasks, return_exceptions=True)
+        else:
+            # Use individual embedding
+            try:
+                embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
+                await asyncio.gather(*embed_tasks, return_exceptions=True)
+            except Exception as e:
+                log_error(f"Error processing document: {e}")
 
         for document in documents:
             cleaned_content = document.content.replace("\x00", "\ufffd")
@@ -196,11 +291,14 @@ class ChromaDb(VectorDb):
 
             metadata["content_hash"] = content_hash
 
+            # Flatten metadata for ChromaDB compatibility
+            flattened_metadata = self._flatten_metadata(metadata)
+
             docs_embeddings.append(document.embedding)
             docs.append(cleaned_content)
             ids.append(doc_id)
-            docs_metadata.append(metadata)
-            log_debug(f"Prepared document: {document.id} | {document.name} | {metadata}")
+            docs_metadata.append(flattened_metadata)
+            log_debug(f"Prepared document: {document.id} | {document.name} | {flattened_metadata}")
 
         if self._collection is None:
             logger.warning("Collection does not exist")
@@ -262,11 +360,14 @@ class ChromaDb(VectorDb):
 
             metadata["content_hash"] = content_hash
 
+            # Flatten metadata for ChromaDB compatibility
+            flattened_metadata = self._flatten_metadata(metadata)
+
             docs_embeddings.append(document.embedding)
             docs.append(cleaned_content)
             ids.append(doc_id)
-            docs_metadata.append(metadata)
-            log_debug(f"Upserted document: {document.id} | {document.name} | {metadata}")
+            docs_metadata.append(flattened_metadata)
+            log_debug(f"Upserted document: {document.id} | {document.name} | {flattened_metadata}")
 
         if self._collection is None:
             logger.warning("Collection does not exist")
@@ -293,8 +394,44 @@ class ChromaDb(VectorDb):
         if not self._collection:
             self._collection = self.client.get_collection(name=self.collection_name)
 
-        embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
-        await asyncio.gather(*embed_tasks, return_exceptions=True)
+        if self.embedder.enable_batch and hasattr(self.embedder, "async_get_embeddings_batch_and_usage"):
+            # Use batch embedding when enabled and supported
+            try:
+                # Extract content from all documents
+                doc_contents = [doc.content for doc in documents]
+
+                # Get batch embeddings and usage
+                embeddings, usages = await self.embedder.async_get_embeddings_batch_and_usage(doc_contents)
+
+                # Process documents with pre-computed embeddings
+                for j, doc in enumerate(documents):
+                    try:
+                        if j < len(embeddings):
+                            doc.embedding = embeddings[j]
+                            doc.usage = usages[j] if j < len(usages) else None
+                    except Exception as e:
+                        logger.error(f"Error assigning batch embedding to document '{doc.name}': {e}")
+
+            except Exception as e:
+                # Check if this is a rate limit error - don't fall back as it would make things worse
+                error_str = str(e).lower()
+                is_rate_limit = any(
+                    phrase in error_str
+                    for phrase in ["rate limit", "too many requests", "429", "trial key", "api calls / minute"]
+                )
+
+                if is_rate_limit:
+                    logger.error(f"Rate limit detected during batch embedding. {e}")
+                    raise e
+                else:
+                    logger.warning(f"Async batch embedding failed, falling back to individual embeddings: {e}")
+                    # Fall back to individual embedding
+                    embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in documents]
+                    await asyncio.gather(*embed_tasks, return_exceptions=True)
+        else:
+            # Use individual embedding
+            embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
+            await asyncio.gather(*embed_tasks, return_exceptions=True)
 
         for document in documents:
             cleaned_content = document.content.replace("\x00", "\ufffd")
@@ -313,11 +450,14 @@ class ChromaDb(VectorDb):
 
             metadata["content_hash"] = content_hash
 
+            # Flatten metadata for ChromaDB compatibility
+            flattened_metadata = self._flatten_metadata(metadata)
+
             docs_embeddings.append(document.embedding)
             docs.append(cleaned_content)
             ids.append(doc_id)
-            docs_metadata.append(metadata)
-            log_debug(f"Upserted document: {document.id} | {document.name} | {metadata}")
+            docs_metadata.append(flattened_metadata)
+            log_debug(f"Upserted document: {document.id} | {document.name} | {flattened_metadata}")
 
         if self._collection is None:
             logger.warning("Collection does not exist")
@@ -338,13 +478,15 @@ class ChromaDb(VectorDb):
             logger.error(f"Error upserting documents by content hash: {e}")
             raise
 
-    def search(self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None) -> List[Document]:
+    def search(
+        self, query: str, limit: int = 5, filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None
+    ) -> List[Document]:
         """Search the collection for a query.
 
         Args:
             query (str): Query to search for.
             limit (int): Number of results to return.
-            filters (Optional[Dict[str, Any]]): Filters to apply while searching.
+            filters (Optional[Union[Dict[str, Any], List[FilterExpr]]]): Filters to apply while searching.
                 Supports ChromaDB's filtering operators:
                 - $eq, $ne: Equality/Inequality
                 - $gt, $gte, $lt, $lte: Numeric comparisons
@@ -353,6 +495,9 @@ class ChromaDb(VectorDb):
         Returns:
             List[Document]: List of search results.
         """
+        if isinstance(filters, list):
+            log_warning("Filter Expressions are not yet supported in ChromaDB. No filters will be applied.")
+            filters = None
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             logger.error(f"Error getting embedding for Query: {query}")
@@ -374,11 +519,11 @@ class ChromaDb(VectorDb):
         # Build search results
         search_results: List[Document] = []
 
-        ids_list = result.get("ids", [[]])
-        metadata_list = result.get("metadatas", [[{}]])
-        documents_list = result.get("documents", [[]])
-        embeddings_list = result.get("embeddings")
-        distances_list = result.get("distances", [[]])
+        ids_list = result.get("ids", [[]])  # type: ignore
+        metadata_list = result.get("metadatas", [[{}]])  # type: ignore
+        documents_list = result.get("documents", [[]])  # type: ignore
+        embeddings_list = result.get("embeddings")  # type: ignore
+        distances_list = result.get("distances", [[]])  # type: ignore
 
         if not ids_list or not metadata_list or not documents_list or embeddings_list is None or not distances_list:
             return search_results
@@ -467,7 +612,7 @@ class ChromaDb(VectorDb):
         return converted
 
     async def async_search(
-        self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None
+        self, query: str, limit: int = 5, filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None
     ) -> List[Document]:
         """Search asynchronously by running in a thread."""
         return await asyncio.to_thread(self.search, query, limit, filters)
@@ -658,7 +803,6 @@ class ChromaDb(VectorDb):
 
         try:
             collection: Collection = self.client.get_collection(name=self.collection_name)
-            print("COLLECTION_----------", collection)
             # Try to get the document by ID
             result = collection.get(ids=[id])
             found_ids = result.get("ids", [])
@@ -747,6 +891,9 @@ class ChromaDb(VectorDb):
                     logger.debug(f"No documents found with content_id: {content_id}")
                     return
 
+                # Flatten the new metadata first
+                flattened_new_metadata = self._flatten_metadata(metadata)
+
                 # Merge metadata for each document
                 updated_metadatas = []
                 for i, current_meta in enumerate(current_metadatas or []):
@@ -754,21 +901,14 @@ class ChromaDb(VectorDb):
                         meta_dict: Dict[str, Any] = {}
                     else:
                         meta_dict = dict(current_meta)  # Convert Mapping to dict
-                    updated_meta: Dict[str, Any] = meta_dict.copy()
-                    updated_meta.update(metadata)
 
-                    if "filters" not in updated_meta:
-                        updated_meta["filters"] = {}
-                    if isinstance(updated_meta["filters"], dict):
-                        updated_meta["filters"].update(metadata)
-                    else:
-                        updated_meta["filters"] = metadata
-                    updated_metadatas.append(updated_meta)
+                    # Update with flattened metadata
+                    meta_dict.update(flattened_new_metadata)
+                    updated_metadatas.append(meta_dict)
 
-                # Update the documents
                 # Convert to the expected type for ChromaDB
-                chroma_metadatas = cast(List[Mapping[str, Union[str, int, float, bool, None]]], updated_metadatas)
-                collection.update(ids=ids, metadatas=chroma_metadatas)
+                chroma_metadatas = cast(List[Mapping[str, Union[str, int, float, bool]]], updated_metadatas)
+                collection.update(ids=ids, metadatas=chroma_metadatas)  # type: ignore
                 logger.debug(f"Updated metadata for {len(ids)} documents with content_id: {content_id}")
 
             except TypeError as te:
@@ -783,3 +923,7 @@ class ChromaDb(VectorDb):
         except Exception as e:
             logger.error(f"Error updating metadata for content_id '{content_id}': {e}")
             raise
+
+    def get_supported_search_types(self) -> List[str]:
+        """Get the supported search types for this vector database."""
+        return []  # ChromaDb doesn't use SearchType enum

agno/vectordb/clickhouse/clickhousedb.py

@@ -1,6 +1,6 @@
 import asyncio
 from hashlib import md5
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from agno.vectordb.clickhouse.index import HNSW
 
@@ -11,9 +11,10 @@ try:
 except ImportError:
     raise ImportError("`clickhouse-connect` not installed. Use `pip install clickhouse-connect` to install it")
 
+from agno.filters import FilterExpr
 from agno.knowledge.document import Document
 from agno.knowledge.embedder import Embedder
-from agno.utils.log import log_debug, log_info, logger
+from agno.utils.log import log_debug, log_info, log_warning, logger
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 
@@ -23,6 +24,8 @@ class Clickhouse(VectorDb):
         self,
         table_name: str,
         host: str,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
         username: Optional[str] = None,
         password: str = "",
         port: int = 0,
@@ -41,9 +44,11 @@ class Clickhouse(VectorDb):
         self.password = password
         self.port = port
         self.dsn = dsn
+        # Initialize base class with name and description
+        super().__init__(name=name, description=description)
+
         self.compress = compress
         self.database_name = database_name
-
         if not client:
             client = clickhouse_connect.get_client(
                 host=self.host,
@@ -81,6 +86,7 @@ class Clickhouse(VectorDb):
         if self.async_client is None:
             self.async_client = await clickhouse_connect.get_async_client(
                 host=self.host,
+                username=self.username,  # type: ignore
                 password=self.password,
                 database=self.database_name,
                 port=self.port,
@@ -228,7 +234,7 @@ class Clickhouse(VectorDb):
             "SELECT name FROM {database_name:Identifier}.{table_name:Identifier} WHERE name = {name:String}",
             parameters=parameters,
         )
-        return bool(result)
+        return len(result.result_rows) > 0 if result.result_rows else False
 
     async def async_name_exists(self, name: str) -> bool:
         """Check if a document with given name exists asynchronously."""
@@ -241,7 +247,7 @@ class Clickhouse(VectorDb):
             "SELECT name FROM {database_name:Identifier}.{table_name:Identifier} WHERE name = {name:String}",
             parameters=parameters,
         )
-        return bool(result)
+        return len(result.result_rows) > 0 if result.result_rows else False
 
     def id_exists(self, id: str) -> bool:
         """
@@ -257,7 +263,7 @@ class Clickhouse(VectorDb):
             "SELECT id FROM {database_name:Identifier}.{table_name:Identifier} WHERE id = {id:String}",
             parameters=parameters,
         )
-        return bool(result)
+        return len(result.result_rows) > 0 if result.result_rows else False
 
     def insert(
         self,
@@ -308,8 +314,44 @@ class Clickhouse(VectorDb):
         rows: List[List[Any]] = []
         async_client = await self._ensure_async_client()
 
-        embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
-        await asyncio.gather(*embed_tasks, return_exceptions=True)
+        if self.embedder.enable_batch and hasattr(self.embedder, "async_get_embeddings_batch_and_usage"):
+            # Use batch embedding when enabled and supported
+            try:
+                # Extract content from all documents
+                doc_contents = [doc.content for doc in documents]
+
+                # Get batch embeddings and usage
+                embeddings, usages = await self.embedder.async_get_embeddings_batch_and_usage(doc_contents)
+
+                # Process documents with pre-computed embeddings
+                for j, doc in enumerate(documents):
+                    try:
+                        if j < len(embeddings):
+                            doc.embedding = embeddings[j]
+                            doc.usage = usages[j] if j < len(usages) else None
+                    except Exception as e:
+                        logger.error(f"Error assigning batch embedding to document '{doc.name}': {e}")
+
+            except Exception as e:
+                # Check if this is a rate limit error - don't fall back as it would make things worse
+                error_str = str(e).lower()
+                is_rate_limit = any(
+                    phrase in error_str
+                    for phrase in ["rate limit", "too many requests", "429", "trial key", "api calls / minute"]
+                )
+
+                if is_rate_limit:
+                    logger.error(f"Rate limit detected during batch embedding. {e}")
+                    raise e
+                else:
+                    logger.warning(f"Async batch embedding failed, falling back to individual embeddings: {e}")
+                    # Fall back to individual embedding
+                    embed_tasks = [doc.async_embed(embedder=self.embedder) for doc in documents]
+                    await asyncio.gather(*embed_tasks, return_exceptions=True)
+        else:
+            # Use individual embedding
+            embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
+            await asyncio.gather(*embed_tasks, return_exceptions=True)
 
         for document in documents:
             cleaned_content = document.content.replace("\x00", "\ufffd")
@@ -407,7 +449,11 @@ class Clickhouse(VectorDb):
             parameters=parameters,
         )
 
-    def search(self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None) -> List[Document]:
+    def search(
+        self, query: str, limit: int = 5, filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None
+    ) -> List[Document]:
+        if filters is not None:
+            log_warning("Filters are not yet supported in Clickhouse. No filters will be applied.")
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             logger.error(f"Error getting embedding for Query: {query}")
@@ -461,11 +507,14 @@ class Clickhouse(VectorDb):
         return search_results
 
     async def async_search(
-        self, query: str, limit: int = 5, filters: Optional[Dict[str, Any]] = None
+        self, query: str, limit: int = 5, filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None
     ) -> List[Document]:
         """Search for documents asynchronously."""
         async_client = await self._ensure_async_client()
 
+        if filters is not None:
+            log_warning("Filters are not yet supported in Clickhouse. No filters will be applied.")
+
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             logger.error(f"Error getting embedding for Query: {query}")
@@ -696,7 +745,7 @@ class Clickhouse(VectorDb):
             "SELECT content_hash FROM {database_name:Identifier}.{table_name:Identifier} WHERE content_hash = {content_hash:String}",
             parameters=parameters,
         )
-        return bool(result)
+        return len(result.result_rows) > 0 if result.result_rows else False
 
     def _delete_by_content_hash(self, content_hash: str) -> bool:
         """
@@ -780,3 +829,7 @@ class Clickhouse(VectorDb):
         except Exception as e:
             logger.error(f"Error updating metadata for content_id '{content_id}': {e}")
             raise
+
+    def get_supported_search_types(self) -> List[str]:
+        """Get the supported search types for this vector database."""
+        return []  # Clickhouse doesn't use SearchType enum