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

agno/vectordb/singlestore/singlestore.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 from hashlib import md5
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 try:
     from sqlalchemy.dialects import mysql
@@ -14,10 +14,11 @@ try:
 except ImportError:
     raise ImportError("`sqlalchemy` not installed")
 
+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
+from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 
@@ -32,6 +33,8 @@ class SingleStore(VectorDb):
         embedder: Optional[Embedder] = None,
         distance: Distance = Distance.cosine,
         reranker: Optional[Reranker] = None,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
         # index: Optional[Union[Ivfflat, HNSW]] = HNSW(),
     ):
         _engine: Optional[Engine] = db_engine
@@ -44,9 +47,11 @@ class SingleStore(VectorDb):
         self.collection: str = collection
         self.schema: Optional[str] = schema
         self.db_url: Optional[str] = db_url
+        # Initialize base class with name and description
+        super().__init__(name=name, description=description)
+
         self.db_engine: Engine = _engine
         self.metadata: MetaData = MetaData(schema=self.schema)
-
         if embedder is None:
             from agno.knowledge.embedder.openai import OpenAIEmbedder
 
@@ -278,7 +283,9 @@ class SingleStore(VectorDb):
             sess.commit()
             log_debug(f"Committed {counter} documents")
 
-    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 for documents based on a query and optional filters.
 
@@ -290,6 +297,8 @@ class SingleStore(VectorDb):
         Returns:
             List[Document]: List of documents that match the query.
         """
+        if filters is not None:
+            log_warning("Filters are not supported in SingleStore. No filters will be applied.")
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             log_error(f"Error getting embedding for Query: {query}")
@@ -428,9 +437,9 @@ class SingleStore(VectorDb):
         try:
             with self.Session.begin() as sess:
                 stmt = delete(self.table).where(self.table.c.id == id)
-                result = sess.execute(stmt)
-                log_info(f"Deleted {result.rowcount} records with ID {id} from table '{self.table.name}'.")
-                return result.rowcount > 0
+                result = sess.execute(stmt)  # type: ignore
+                log_info(f"Deleted {result.rowcount} records with ID {id} from table '{self.table.name}'.")  # type: ignore
+                return result.rowcount > 0  # type: ignore
         except Exception as e:
             log_error(f"Error deleting document with ID {id}: {e}")
             return False
@@ -444,11 +453,11 @@ class SingleStore(VectorDb):
         try:
             with self.Session.begin() as sess:
                 stmt = delete(self.table).where(self.table.c.content_id == content_id)
-                result = sess.execute(stmt)
+                result = sess.execute(stmt)  # type: ignore
                 log_info(
-                    f"Deleted {result.rowcount} records with content_id {content_id} from table '{self.table.name}'."
+                    f"Deleted {result.rowcount} records with content_id {content_id} from table '{self.table.name}'."  # type: ignore
                 )
-                return result.rowcount > 0
+                return result.rowcount > 0  # type: ignore
         except Exception as e:
             log_error(f"Error deleting document with content_id {content_id}: {e}")
             return False
@@ -462,9 +471,9 @@ class SingleStore(VectorDb):
         try:
             with self.Session.begin() as sess:
                 stmt = delete(self.table).where(self.table.c.name == name)
-                result = sess.execute(stmt)
-                log_info(f"Deleted {result.rowcount} records with name '{name}' from table '{self.table.name}'.")
-                return result.rowcount > 0
+                result = sess.execute(stmt)  # type: ignore
+                log_info(f"Deleted {result.rowcount} records with name '{name}' from table '{self.table.name}'.")  # type: ignore
+                return result.rowcount > 0  # type: ignore
         except Exception as e:
             log_error(f"Error deleting document with name {name}: {e}")
             return False
@@ -480,9 +489,9 @@ class SingleStore(VectorDb):
                 # Convert metadata to JSON string for comparison
                 metadata_json = json.dumps(metadata, sort_keys=True)
                 stmt = delete(self.table).where(self.table.c.meta_data == metadata_json)
-                result = sess.execute(stmt)
-                log_info(f"Deleted {result.rowcount} records with metadata {metadata} from table '{self.table.name}'.")
-                return result.rowcount > 0
+                result = sess.execute(stmt)  # type: ignore
+                log_info(f"Deleted {result.rowcount} records with metadata {metadata} from table '{self.table.name}'.")  # type: ignore
+                return result.rowcount > 0  # type: ignore
         except Exception as e:
             log_error(f"Error deleting documents with metadata {metadata}: {e}")
             return False
@@ -496,8 +505,44 @@ class SingleStore(VectorDb):
         documents: List[Document],
         filters: Optional[Dict[str, Any]] = None,
     ) -> None:
-        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:
+                        log_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:
+                    log_error(f"Rate limit detected during batch embedding. {e}")
+                    raise e
+                else:
+                    log_error(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)
 
         with self.Session.begin() as sess:
             counter = 0
@@ -543,8 +588,45 @@ class SingleStore(VectorDb):
             filters (Optional[Dict[str, Any]]): Optional filters for the upsert.
             batch_size (int): Number of documents to upsert in each batch.
         """
-        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:
+                        log_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:
+                    log_error(f"Rate limit detected during batch embedding. {e}")
+                    raise e
+                else:
+                    log_error(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)
 
         with self.Session.begin() as sess:
             counter = 0
@@ -588,7 +670,7 @@ class SingleStore(VectorDb):
             log_debug(f"Committed {counter} documents")
 
     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]:
         return self.search(query=query, limit=limit, filters=filters)
 
@@ -616,11 +698,11 @@ class SingleStore(VectorDb):
         try:
             with self.Session.begin() as sess:
                 stmt = delete(self.table).where(self.table.c.content_hash == content_hash)
-                result = sess.execute(stmt)
+                result = sess.execute(stmt)  # type: ignore
                 log_info(
-                    f"Deleted {result.rowcount} records with content_hash '{content_hash}' from table '{self.table.name}'."
+                    f"Deleted {result.rowcount} records with content_hash '{content_hash}' from table '{self.table.name}'."  # type: ignore
                 )
-                return result.rowcount > 0
+                return result.rowcount > 0  # type: ignore
         except Exception as e:
             log_error(f"Error deleting documents with content_hash {content_hash}: {e}")
             return False
@@ -639,7 +721,7 @@ class SingleStore(VectorDb):
             with self.Session.begin() as sess:
                 # Find documents with the given content_id
                 stmt = select(self.table).where(self.table.c.content_id == content_id)
-                result = sess.execute(stmt)
+                result = sess.execute(stmt)  # type: ignore
 
                 updated_count = 0
                 for row in result:
@@ -675,3 +757,7 @@ class SingleStore(VectorDb):
         except Exception as e:
             log_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 []  # SingleStore doesn't use SearchType enum

agno/vectordb/surrealdb/surrealdb.py

@@ -11,9 +11,10 @@ except ImportError as e:
     msg = "The `surrealdb` package is not installed. Please install it via `pip install surrealdb`."
     raise ImportError(msg) from e
 
+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_error, log_info
+from agno.utils.log import log_debug, log_error, log_info, log_warning
 from agno.vectordb.base import VectorDb
 from agno.vectordb.distance import Distance
 
@@ -107,11 +108,13 @@ class SurrealDb(VectorDb):
         m: int = 12,
         search_ef: int = 40,
         embedder: Optional[Embedder] = None,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        id: Optional[str] = None,
     ):
         """Initialize SurrealDB connection.
 
         Args:
-            url: SurrealDB server URL (e.g. ws://localhost:8000/rpc)
             client: A blocking connection, either HTTP or WS
             async_client: An async connection, either HTTP or WS (default: None)
             collection: Collection name to store documents (default: documents)
@@ -122,6 +125,17 @@ class SurrealDb(VectorDb):
             embedder: Embedder instance for creating embeddings (default: OpenAIEmbedder)
 
         """
+        # Dynamic ID generation based on unique identifiers
+        if id is None:
+            from agno.utils.string import generate_id
+
+            client_info = str(client) if client else str(async_client) if async_client else "default"
+            seed = f"{client_info}#{collection}"
+            id = generate_id(seed)
+
+        # Initialize base class with name, description, and generated ID
+        super().__init__(id=id, name=name, description=description)
+
         # Embedder for embedding the document contents
         if embedder is None:
             from agno.knowledge.embedder.openai import OpenAIEmbedder
@@ -131,7 +145,6 @@ class SurrealDb(VectorDb):
         self.embedder: Embedder = embedder
         self.dimensions = self.embedder.dimensions
         self.collection = collection
-
         # Convert Distance enum to SurrealDB distance type
         self.distance = {Distance.cosine: "COSINE", Distance.l2: "EUCLIDEAN", Distance.max_inner_product: "DOT"}[
             distance
@@ -306,7 +319,9 @@ class SurrealDb(VectorDb):
             thing = f"{self.collection}:{doc.id}" if doc.id else self.collection
             self.client.query(self.UPSERT_QUERY.format(thing=thing), data)
 
-    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 for similar documents.
 
         Args:
@@ -318,6 +333,9 @@ class SurrealDb(VectorDb):
             A list of documents that are similar to the query.
 
         """
+        if isinstance(filters, List):
+            log_warning("Filters Expressions are not supported in SurrealDB. No filters will be applied.")
+            filters = None
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             log_error(f"Error getting embedding for Query: {query}")
@@ -548,7 +566,7 @@ class SurrealDb(VectorDb):
         self,
         query: str,
         limit: int = 5,
-        filters: Optional[Dict[str, Any]] = None,
+        filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None,
     ) -> List[Document]:
         """Search for similar documents asynchronously.
 
@@ -561,6 +579,10 @@ class SurrealDb(VectorDb):
             A list of documents that are similar to the query.
 
         """
+        if isinstance(filters, List):
+            log_warning("Filters Expressions are not supported in SurrealDB. No filters will be applied.")
+            filters = None
+
         query_embedding = self.embedder.get_embedding(query)
         if query_embedding is None:
             log_error(f"Error getting embedding for Query: {query}")
@@ -671,3 +693,7 @@ class SurrealDb(VectorDb):
         except Exception as e:
             log_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 []  # SurrealDb doesn't use SearchType enum

agno/vectordb/upstashdb/upstashdb.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 try:
     from upstash_vector import Index, Vector
@@ -9,10 +9,11 @@ except ImportError:
         "The `upstash-vector` package is not installed, please install using `pip install upstash-vector`"
     )
 
+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_info, logger
+from agno.utils.log import log_info, log_warning, logger
 from agno.vectordb.base import VectorDb
 
 DEFAULT_NAMESPACE = ""
@@ -32,6 +33,8 @@ class UpstashVectorDb(VectorDb):
         embedder (Optional[Embedder], optional): The embedder to use. If None, uses Upstash hosted embedding models.
         namespace (Optional[str], optional): The namespace to use. Defaults to DEFAULT_NAMESPACE.
         reranker (Optional[Reranker], optional): The reranker to use. Defaults to None.
+        name (Optional[str], optional): The name of the vector database. Defaults to None.
+        description (Optional[str], optional): The description of the vector database. Defaults to None.
         **kwargs: Additional keyword arguments.
     """
 
@@ -45,8 +48,28 @@ class UpstashVectorDb(VectorDb):
         embedder: Optional[Embedder] = None,
         namespace: Optional[str] = DEFAULT_NAMESPACE,
         reranker: Optional[Reranker] = None,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        id: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
+        # Validate required parameters
+        if not url:
+            raise ValueError("URL must be provided.")
+        if not token:
+            raise ValueError("Token must be provided.")
+
+        # Dynamic ID generation based on unique identifiers
+        if id is None:
+            from agno.utils.string import generate_id
+
+            namespace_identifier = namespace or DEFAULT_NAMESPACE
+            seed = f"{url}#{namespace_identifier}"
+            id = generate_id(seed)
+
+        # Initialize base class with name, description, and generated ID
+        super().__init__(id=id, name=name, description=description)
+
         self._index: Optional[Index] = None
         self.url: str = url
         self.token: str = token
@@ -56,7 +79,6 @@ class UpstashVectorDb(VectorDb):
         self.namespace: str = namespace if namespace is not None else DEFAULT_NAMESPACE
         self.kwargs: Dict[str, Any] = kwargs
         self.use_upstash_embeddings: bool = embedder is None
-
         if embedder is None:
             logger.warning(
                 "You have not provided an embedder, using Upstash hosted embedding models. "
@@ -303,7 +325,7 @@ class UpstashVectorDb(VectorDb):
         self,
         query: str,
         limit: int = 5,
-        filters: Optional[Dict[str, Any]] = None,
+        filters: Optional[Union[Dict[str, Any], List[FilterExpr]]] = None,
         namespace: Optional[str] = None,
     ) -> List[Document]:
         """Search for documents in the index.
@@ -316,7 +338,9 @@ class UpstashVectorDb(VectorDb):
             List[Document]: List of matching documents.
         """
         _namespace = self.namespace if namespace is None else namespace
-
+        if isinstance(filters, List):
+            log_warning("Filters Expressions are not supported in UpstashDB. No filters will be applied.")
+            filters = None
         filter_str = "" if filters is None else str(filters)
 
         if not self.use_upstash_embeddings and self.embedder is not None:
@@ -504,8 +528,48 @@ class UpstashVectorDb(VectorDb):
         _namespace = self.namespace if namespace is None else namespace
         vectors = []
 
-        embed_tasks = [document.async_embed(embedder=self.embedder) for document in documents]
-        await asyncio.gather(*embed_tasks, return_exceptions=True)
+        if (
+            self.embedder
+            and 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 i, document in enumerate(documents):
             if document.id is None:
@@ -562,7 +626,7 @@ class UpstashVectorDb(VectorDb):
         self.index.upsert(vectors, namespace=_namespace)
 
     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]:
         raise NotImplementedError(f"Async not supported on {self.__class__.__name__}.")
 
@@ -648,3 +712,7 @@ class UpstashVectorDb(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 []  # UpstashVectorDb doesn't use SearchType enum