agno 2.4.6__py3-none-any.whl → 2.4.7__py3-none-any.whl

agno/knowledge/reranker/aws_bedrock.py

@@ -0,0 +1,299 @@
+from os import getenv
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import ConfigDict, Field
+
+from agno.knowledge.document import Document
+from agno.knowledge.reranker.base import Reranker
+from agno.utils.log import logger
+
+try:
+    from boto3 import client as AwsClient
+    from boto3.session import Session
+    from botocore.exceptions import ClientError
+except ImportError:
+    raise ImportError("`boto3` not installed. Please install it via `pip install boto3`.")
+
+
+# Model ID constants
+AMAZON_RERANK_V1 = "amazon.rerank-v1:0"
+COHERE_RERANK_V3_5 = "cohere.rerank-v3-5:0"
+
+# Type alias for supported models
+RerankerModel = Literal["amazon.rerank-v1:0", "cohere.rerank-v3-5:0"]
+
+
+class AwsBedrockReranker(Reranker):
+    """
+    AWS Bedrock reranker supporting Amazon Rerank 1.0 and Cohere Rerank 3.5 models.
+
+    This reranker uses the unified Bedrock Rerank API (bedrock-agent-runtime)
+    which provides a consistent interface for both model providers.
+
+    To use this reranker, you need to either:
+    1. Set the following environment variables:
+       - AWS_ACCESS_KEY_ID
+       - AWS_SECRET_ACCESS_KEY
+       - AWS_REGION
+    2. Or provide a boto3 Session object
+
+    Args:
+        model (str): The model ID to use. Options:
+            - 'amazon.rerank-v1:0' (Amazon Rerank 1.0)
+            - 'cohere.rerank-v3-5:0' (Cohere Rerank 3.5)
+            Default is 'cohere.rerank-v3-5:0'.
+        top_n (Optional[int]): Number of top results to return after reranking.
+            If None, returns all documents reranked.
+        aws_region (Optional[str]): The AWS region to use.
+        aws_access_key_id (Optional[str]): The AWS access key ID to use.
+        aws_secret_access_key (Optional[str]): The AWS secret access key to use.
+        session (Optional[Session]): A boto3 Session object for authentication.
+        additional_model_request_fields (Optional[Dict]): Additional model-specific
+            parameters to pass in the request (e.g., Cohere-specific options).
+
+    Example:
+        ```python
+        from agno.knowledge.reranker.aws_bedrock import AwsBedrockReranker
+
+        # Using Cohere Rerank 3.5 (default)
+        reranker = AwsBedrockReranker(
+            model="cohere.rerank-v3-5:0",
+            top_n=5,
+            aws_region="us-west-2",
+        )
+
+        # Using Amazon Rerank 1.0
+        reranker = AwsBedrockReranker(
+            model="amazon.rerank-v1:0",
+            top_n=10,
+            aws_region="us-west-2",
+        )
+
+        # Rerank documents
+        reranked_docs = reranker.rerank(query="What is machine learning?", documents=docs)
+        ```
+
+    Note:
+        - Amazon Rerank 1.0 is NOT available in us-east-1 (N. Virginia).
+          Use Cohere Rerank 3.5 in that region.
+        - Maximum 1000 documents per request.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)
+
+    model: str = Field(default=COHERE_RERANK_V3_5, description="Reranker model ID")
+    top_n: Optional[int] = Field(default=None, description="Number of top results to return")
+
+    aws_region: Optional[str] = Field(default=None, description="AWS region")
+    aws_access_key_id: Optional[str] = Field(default=None, description="AWS access key ID")
+    aws_secret_access_key: Optional[str] = Field(default=None, description="AWS secret access key")
+    session: Optional[Session] = Field(default=None, description="Boto3 session", exclude=True)
+
+    additional_model_request_fields: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Additional model-specific request parameters",
+    )
+
+    _client: Optional[AwsClient] = None
+
+    @property
+    def client(self) -> AwsClient:
+        """
+        Returns a bedrock-agent-runtime client for the Rerank API.
+
+        Returns:
+            AwsClient: An instance of the bedrock-agent-runtime client.
+        """
+        if self._client is not None:
+            return self._client
+
+        if self.session:
+            self._client = self.session.client("bedrock-agent-runtime")
+            return self._client
+
+        aws_access_key_id = self.aws_access_key_id or getenv("AWS_ACCESS_KEY_ID")
+        aws_secret_access_key = self.aws_secret_access_key or getenv("AWS_SECRET_ACCESS_KEY")
+        aws_region = self.aws_region or getenv("AWS_REGION")
+
+        if not aws_access_key_id or not aws_secret_access_key:
+            # Fall back to default credential chain
+            self._client = AwsClient(
+                service_name="bedrock-agent-runtime",
+                region_name=aws_region,
+            )
+        else:
+            self._client = AwsClient(
+                service_name="bedrock-agent-runtime",
+                region_name=aws_region,
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+            )
+
+        return self._client
+
+    def _get_model_arn(self) -> str:
+        """
+        Constructs the full model ARN for the reranker model.
+
+        Returns:
+            str: The model ARN.
+        """
+        region = self.aws_region or getenv("AWS_REGION", "us-west-2")
+        return f"arn:aws:bedrock:{region}::foundation-model/{self.model}"
+
+    def _build_sources(self, documents: List[Document]) -> List[Dict[str, Any]]:
+        """
+        Convert Document objects to Bedrock Rerank API source format.
+
+        Args:
+            documents: List of Document objects to convert.
+
+        Returns:
+            List of RerankSource objects for the API.
+        """
+        sources = []
+        for doc in documents:
+            # Use text format for document content
+            source = {
+                "type": "INLINE",
+                "inlineDocumentSource": {
+                    "type": "TEXT",
+                    "textDocument": {
+                        "text": doc.content,
+                    },
+                },
+            }
+            sources.append(source)
+        return sources
+
+    def _rerank(self, query: str, documents: List[Document]) -> List[Document]:
+        """
+        Internal method to perform reranking via Bedrock Rerank API.
+
+        Args:
+            query: The query string to rank documents against.
+            documents: List of Document objects to rerank.
+
+        Returns:
+            List of Document objects sorted by relevance score.
+        """
+        if not documents:
+            return []
+
+        # Validate top_n
+        top_n = self.top_n
+        if top_n is not None and top_n <= 0:
+            logger.warning(f"top_n should be a positive integer, got {self.top_n}, setting top_n to None")
+            top_n = None
+
+        # Build the request
+        rerank_request: Dict[str, Any] = {
+            "queries": [
+                {
+                    "type": "TEXT",
+                    "textQuery": {
+                        "text": query,
+                    },
+                }
+            ],
+            "sources": self._build_sources(documents),
+            "rerankingConfiguration": {
+                "type": "BEDROCK_RERANKING_MODEL",
+                "bedrockRerankingConfiguration": {
+                    "modelConfiguration": {
+                        "modelArn": self._get_model_arn(),
+                    },
+                },
+            },
+        }
+
+        # Add numberOfResults if top_n is specified
+        if top_n is not None:
+            rerank_request["rerankingConfiguration"]["bedrockRerankingConfiguration"]["numberOfResults"] = top_n
+
+        # Add additional model request fields if provided
+        if self.additional_model_request_fields:
+            rerank_request["rerankingConfiguration"]["bedrockRerankingConfiguration"]["modelConfiguration"][
+                "additionalModelRequestFields"
+            ] = self.additional_model_request_fields
+
+        # Call the Rerank API
+        response = self.client.rerank(**rerank_request)
+
+        # Process results
+        reranked_docs: List[Document] = []
+        results = response.get("results", [])
+
+        for result in results:
+            index = result.get("index")
+            relevance_score = result.get("relevanceScore")
+
+            if index is not None and index < len(documents):
+                doc = documents[index]
+                doc.reranking_score = relevance_score
+                reranked_docs.append(doc)
+
+        # Results from API are already sorted by relevance, but ensure sorting
+        reranked_docs.sort(
+            key=lambda x: x.reranking_score if x.reranking_score is not None else float("-inf"),
+            reverse=True,
+        )
+
+        return reranked_docs
+
+    def rerank(self, query: str, documents: List[Document]) -> List[Document]:
+        """
+        Rerank documents based on their relevance to a query.
+
+        Args:
+            query: The query string to rank documents against.
+            documents: List of Document objects to rerank.
+
+        Returns:
+            List of Document objects sorted by relevance score (highest first).
+            Each document will have its `reranking_score` attribute set.
+        """
+        try:
+            return self._rerank(query=query, documents=documents)
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code", "Unknown")
+            error_message = e.response.get("Error", {}).get("Message", str(e))
+            logger.error(f"AWS Bedrock Rerank API error ({error_code}): {error_message}. Returning original documents.")
+            return documents
+        except Exception as e:
+            logger.error(f"Error reranking documents: {e}. Returning original documents.")
+            return documents
+
+
+class CohereBedrockReranker(AwsBedrockReranker):
+    """
+    Convenience class for Cohere Rerank 3.5 on AWS Bedrock.
+
+    This is a pre-configured AwsBedrockReranker using the Cohere Rerank 3.5 model.
+
+    Example:
+        ```python
+        reranker = CohereBedrockReranker(top_n=5, aws_region="us-west-2")
+        reranked_docs = reranker.rerank(query="What is AI?", documents=docs)
+        ```
+    """
+
+    model: str = Field(default=COHERE_RERANK_V3_5)
+
+
+class AmazonReranker(AwsBedrockReranker):
+    """
+    Convenience class for Amazon Rerank 1.0 on AWS Bedrock.
+
+    This is a pre-configured AwsBedrockReranker using the Amazon Rerank 1.0 model.
+
+    Note: Amazon Rerank 1.0 is NOT available in us-east-1 (N. Virginia).
+
+    Example:
+        ```python
+        reranker = AmazonReranker(top_n=5, aws_region="us-west-2")
+        reranked_docs = reranker.rerank(query="What is AI?", documents=docs)
+        ```
+    """
+
+    model: str = Field(default=AMAZON_RERANK_V1)

agno/learn/machine.py

@@ -645,12 +645,11 @@ class LearningMachine:
         for name, store in self.stores.items():
             try:
                 result = await store.arecall(**context)
-                if result is not None:
-                    results[name] = result
-                    try:
-                        log_debug(f"Recalled from {name}: {result}")
-                    except Exception:
-                        pass
+                results[name] = result
+                try:
+                    log_debug(f"Recalled from {name}: {result}")
+                except Exception:
+                    pass
             except Exception as e:
                 log_warning(f"Error recalling from {name}: {e}")
 

agno/run/workflow.py

@@ -311,6 +311,9 @@ class ConditionExecutionCompletedEvent(BaseWorkflowRunOutputEvent):
     condition_result: Optional[bool] = None
     executed_steps: Optional[int] = None
 
+    # Which branch was executed: "if", "else", or None (condition false with no else_steps)
+    branch: Optional[str] = None
+
     # Results from executed steps
     step_results: List[StepOutput] = field(default_factory=list)
 

agno/tools/mcp/mcp.py

@@ -74,6 +74,13 @@ class MCPTools(Toolkit):
                 Only relevant with HTTP transports (Streamable HTTP or SSE).
                 Creates a new session per agent run with dynamic headers merged into connection config.
         """
+        # Extract these before super().__init__() to bypass early validation
+        # (tools aren't available until build_tools() is called)
+        requires_confirmation_tools = kwargs.pop("requires_confirmation_tools", None)
+        external_execution_required_tools = kwargs.pop("external_execution_required_tools", None)
+        stop_after_tool_call_tools = kwargs.pop("stop_after_tool_call_tools", None)
+        show_result_tools = kwargs.pop("show_result_tools", None)
+
         super().__init__(name="MCPTools", **kwargs)
 
         if url is not None:
@@ -92,6 +99,10 @@ class MCPTools(Toolkit):
         # because tools are not available until `initialize()` is called.
         self.include_tools = include_tools
         self.exclude_tools = exclude_tools
+        self.requires_confirmation_tools = requires_confirmation_tools or []
+        self.external_execution_required_tools = external_execution_required_tools or []
+        self.stop_after_tool_call_tools = stop_after_tool_call_tools or []
+        self.show_result_tools = show_result_tools or []
         self.refresh_connection = refresh_connection
         self.tool_name_prefix = tool_name_prefix
 
@@ -575,13 +586,27 @@ class MCPTools(Toolkit):
                         mcp_tools_instance=self,
                     )
                     # Create a Function for the tool
+                    # Apply toolkit-level settings
+                    tool_name = tool.name
+                    stop_after = tool_name in self.stop_after_tool_call_tools
+                    show_result = tool_name in self.show_result_tools or stop_after
+
                     f = Function(
-                        name=tool_name_prefix + tool.name,
+                        name=tool_name_prefix + tool_name,
                         description=tool.description,
                         parameters=tool.inputSchema,
                         entrypoint=entrypoint,
                         # Set skip_entrypoint_processing to True to avoid processing the entrypoint
                         skip_entrypoint_processing=True,
+                        # Apply toolkit-level settings for HITL and control flow
+                        requires_confirmation=tool_name in self.requires_confirmation_tools,
+                        external_execution=tool_name in self.external_execution_required_tools,
+                        stop_after_tool_call=stop_after,
+                        show_result=show_result,
+                        # Apply toolkit-level cache settings
+                        cache_results=self.cache_results,
+                        cache_dir=self.cache_dir,
+                        cache_ttl=self.cache_ttl,
                     )
 
                     # Register the Function with the toolkit

agno/vectordb/lancedb/lance_db.py

@@ -104,7 +104,7 @@ class LanceDb(VectorDb):
         self.async_connection: Optional[lancedb.AsyncConnection] = async_connection
         self.async_table: Optional[lancedb.db.AsyncTable] = async_table
 
-        if table_name and table_name in self.connection.table_names():
+        if table_name and table_name in self.connection.list_tables().tables:
             # Open the table if it exists
             try:
                 self.table = self.connection.open_table(name=table_name)
@@ -186,8 +186,8 @@ class LanceDb(VectorDb):
             self.async_connection = await lancedb.connect_async(self.uri)
         # Only try to open table if it exists and we don't have it already
         if self.async_table is None:
-            table_names = await self.async_connection.table_names()
-            if self.table_name in table_names:
+            table_list = await self.async_connection.list_tables()
+            if self.table_name in table_list.tables:
                 try:
                     self.async_table = await self.async_connection.open_table(self.table_name)
                 except ValueError:
@@ -199,7 +199,7 @@ class LanceDb(VectorDb):
         """Refresh the sync connection to see changes made by async operations."""
         try:
             # Re-establish sync connection to see async changes
-            if self.connection and self.table_name in self.connection.table_names():
+            if self.connection is not None and self.table_name in self.connection.list_tables().tables:
                 self.table = self.connection.open_table(self.table_name)
         except Exception as e:
             log_debug(f"Could not refresh sync connection: {e}")
@@ -459,7 +459,7 @@ class LanceDb(VectorDb):
         Returns:
             List[Document]: List of matching documents
         """
-        if self.connection:
+        if self.connection is not None:
             self.table = self.connection.open_table(name=self.table_name)
 
         results = None
@@ -641,8 +641,8 @@ class LanceDb(VectorDb):
         # If we have an async table that was created, the table exists
         if self.async_table is not None:
             return True
-        if self.connection:
-            return self.table_name in self.connection.table_names()
+        if self.connection is not None:
+            return self.table_name in self.connection.list_tables().tables
         return False
 
     async def async_exists(self) -> bool:
@@ -653,8 +653,8 @@ class LanceDb(VectorDb):
         # Check if table exists in database without trying to open it
         if self.async_connection is None:
             self.async_connection = await lancedb.connect_async(self.uri)
-        table_names = await self.async_connection.table_names()
-        return self.table_name in table_names
+        table_list = await self.async_connection.list_tables()
+        return self.table_name in table_list.tables
 
     async def async_get_count(self) -> int:
         """Get the number of rows in the table asynchronously."""