bedrock-agentcore-starter-toolkit 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

bedrock_agentcore_starter_toolkit/services/runtime.py

@@ -0,0 +1,473 @@
+"""BedrockAgentCore service client for agent management."""
+
+import json
+import logging
+import time
+import urllib.parse
+import uuid
+from typing import Any, Dict, Optional
+
+import boto3
+import requests
+from botocore.exceptions import ClientError
+
+from ..utils.endpoints import get_control_plane_endpoint, get_data_plane_endpoint
+
+
+def generate_session_id() -> str:
+    """Generate session ID."""
+    return str(uuid.uuid4())
+
+
+def _handle_http_response(response) -> dict:
+    response.raise_for_status()
+    if "text/event-stream" in response.headers.get("content-type", ""):
+        return _handle_streaming_response(response)
+    else:
+        # Check if response has content
+        if not response.content:
+            raise ValueError("Empty response from agent endpoint")
+
+        return {"response": response.text}
+
+
+def _handle_aws_response(response) -> dict:
+    if "text/event-stream" in response.get("contentType", ""):
+        return _handle_streaming_response(response["response"])
+    else:
+        try:
+            events = []
+            for event in response.get("response", []):
+                events.append(event)
+        except Exception as e:
+            events = [f"Error reading EventStream: {e}"]
+
+        response["response"] = events
+        return response
+
+
+def _handle_streaming_response(response) -> Dict[str, Any]:
+    logger = logging.getLogger("bedrock_agentcore.stream")
+    logger.setLevel(logging.INFO)
+
+    content = []
+    for line in response.iter_lines(chunk_size=1):
+        if line:
+            line = line.decode("utf-8")
+            if line.startswith("data: "):
+                line = line[6:]
+                logger.info(line)
+                content.append(line)
+
+    return {"response": "\n".join(content)}
+
+
+class BedrockAgentCoreClient:
+    """Bedrock AgentCore client for agent management."""
+
+    def __init__(self, region: str):
+        """Initialize Bedrock AgentCore client.
+
+        Args:
+            region: AWS region for the client
+        """
+        self.region = region
+        self.logger = logging.getLogger(f"bedrock_agentcore.runtime.{region}")
+
+        # Get endpoint URLs and log them
+        control_plane_url = get_control_plane_endpoint(region)
+        data_plane_url = get_data_plane_endpoint(region)
+
+        self.logger.debug("Initializing Bedrock AgentCore client for region: %s", region)
+        self.logger.debug("Control plane: %s", control_plane_url)
+        self.logger.debug("Data plane: %s", data_plane_url)
+
+        self.client = boto3.client("bedrock-agentcore-control", region_name=region, endpoint_url=control_plane_url)
+        self.dataplane_client = boto3.client("bedrock-agentcore", region_name=region, endpoint_url=data_plane_url)
+
+    def create_agent(
+        self,
+        agent_name: str,
+        image_uri: str,
+        execution_role_arn: str,
+        network_config: Optional[Dict] = None,
+        authorizer_config: Optional[Dict] = None,
+        protocol_config: Optional[Dict] = None,
+        env_vars: Optional[Dict] = None,
+        auto_update_on_conflict: bool = False,
+    ) -> Dict[str, str]:
+        """Create new agent."""
+        self.logger.info("Creating agent '%s' with image URI: %s", agent_name, image_uri)
+        try:
+            # Build parameters dict, only including optional configs when present
+            params = {
+                "agentRuntimeName": agent_name,
+                "agentRuntimeArtifact": {"containerConfiguration": {"containerUri": image_uri}},
+                "roleArn": execution_role_arn,
+            }
+
+            if network_config is not None:
+                params["networkConfiguration"] = network_config
+
+            if authorizer_config is not None:
+                params["authorizerConfiguration"] = authorizer_config
+
+            if protocol_config is not None:
+                params["protocolConfiguration"] = protocol_config
+
+            if env_vars is not None:
+                params["environmentVariables"] = env_vars
+
+            resp = self.client.create_agent_runtime(**params)
+            agent_id = resp["agentRuntimeId"]
+            agent_arn = resp["agentRuntimeArn"]
+            self.logger.info("Successfully created agent '%s' with ID: %s, ARN: %s", agent_name, agent_id, agent_arn)
+            return {"id": agent_id, "arn": agent_arn}
+        except ClientError as e:
+            error_code = e.response.get("Error", {}).get("Code")
+            if error_code == "ConflictException":
+                if not auto_update_on_conflict:
+                    self.logger.error("Agent '%s' already exists and auto_update_on_conflict is disabled", agent_name)
+                    raise
+
+                self.logger.info("Agent '%s' already exists, searching for existing agent...", agent_name)
+
+                # Find existing agent by name
+                existing_agent = self.find_agent_by_name(agent_name)
+
+                if not existing_agent:
+                    raise RuntimeError(
+                        f"ConflictException occurred but couldn't find existing agent '{agent_name}'. "
+                        f"This might be a permissions issue or the agent name might be different."
+                    ) from e
+
+                # Extract existing agent details
+                existing_agent_id = existing_agent["agentRuntimeId"]
+                existing_agent_arn = existing_agent["agentRuntimeArn"]
+
+                self.logger.info("Found existing agent ID: %s, updating instead...", existing_agent_id)
+
+                # Update the existing agent
+                self.update_agent(
+                    existing_agent_id,
+                    image_uri,
+                    execution_role_arn,
+                    network_config,
+                    authorizer_config,
+                    protocol_config,
+                    env_vars,
+                )
+
+                # Return the existing agent info (keeping the original ID and ARN)
+                return {"id": existing_agent_id, "arn": existing_agent_arn}
+            else:
+                # Re-raise other ClientErrors
+                raise
+        except Exception as e:
+            self.logger.error("Failed to create agent '%s': %s", agent_name, str(e))
+            raise
+
+    def update_agent(
+        self,
+        agent_id: str,
+        image_uri: str,
+        execution_role_arn: str,
+        network_config: Optional[Dict] = None,
+        authorizer_config: Optional[Dict] = None,
+        protocol_config: Optional[Dict] = None,
+        env_vars: Optional[Dict] = None,
+    ) -> Dict[str, str]:
+        """Update existing agent."""
+        self.logger.info("Updating agent ID '%s' with image URI: %s", agent_id, image_uri)
+        try:
+            # Build parameters dict, only including optional configs when present
+            params = {
+                "agentRuntimeId": agent_id,
+                "agentRuntimeArtifact": {"containerConfiguration": {"containerUri": image_uri}},
+                "roleArn": execution_role_arn,
+            }
+
+            if network_config is not None:
+                params["networkConfiguration"] = network_config
+
+            if authorizer_config is not None:
+                params["authorizerConfiguration"] = authorizer_config
+
+            if protocol_config is not None:
+                params["protocolConfiguration"] = protocol_config
+
+            if env_vars is not None:
+                params["environmentVariables"] = env_vars
+
+            resp = self.client.update_agent_runtime(**params)
+            agent_arn = resp["agentRuntimeArn"]
+            self.logger.info("Successfully updated agent ID '%s', ARN: %s", agent_id, agent_arn)
+            return {"id": agent_id, "arn": agent_arn}
+        except Exception as e:
+            self.logger.error("Failed to update agent ID '%s': %s", agent_id, str(e))
+            raise
+
+    def list_agents(self, max_results: int = 100) -> list:
+        """List all agent runtimes, handling pagination."""
+        all_agents = []
+        next_token = None
+
+        try:
+            while True:
+                params = {"maxResults": max_results}
+                if next_token:
+                    params["nextToken"] = next_token
+
+                response = self.client.list_agent_runtimes(**params)
+                agents = response.get("agentRuntimes", [])
+                all_agents.extend(agents)
+
+                next_token = response.get("nextToken")
+                if not next_token:
+                    break
+
+            return all_agents
+        except Exception as e:
+            self.logger.error("Failed to list agents: %s", str(e))
+            raise
+
+    def find_agent_by_name(self, agent_name: str) -> Optional[Dict]:
+        """Find an agent by name, reusing list_agents method."""
+        try:
+            # Get all agents using the existing method
+            all_agents = self.list_agents()
+
+            # Search for the specific agent by name
+            for agent in all_agents:
+                if agent.get("agentRuntimeName") == agent_name:
+                    return agent
+
+            return None  # Agent not found
+        except Exception as e:
+            self.logger.error("Failed to search for agent '%s': %s", agent_name, str(e))
+            raise
+
+    def create_or_update_agent(
+        self,
+        agent_id: Optional[str],
+        agent_name: str,
+        image_uri: str,
+        execution_role_arn: str,
+        network_config: Optional[Dict] = None,
+        authorizer_config: Optional[Dict] = None,
+        protocol_config: Optional[Dict] = None,
+        env_vars: Optional[Dict] = None,
+        auto_update_on_conflict: bool = False,
+    ) -> Dict[str, str]:
+        """Create or update agent."""
+        if agent_id:
+            return self.update_agent(
+                agent_id, image_uri, execution_role_arn, network_config, authorizer_config, protocol_config, env_vars
+            )
+        return self.create_agent(
+            agent_name,
+            image_uri,
+            execution_role_arn,
+            network_config,
+            authorizer_config,
+            protocol_config,
+            env_vars,
+            auto_update_on_conflict,
+        )
+
+    def wait_for_agent_endpoint_ready(self, agent_id: str, endpoint_name: str = "DEFAULT", max_wait: int = 120) -> str:
+        """Wait for agent endpoint to be ready.
+
+        Args:
+            agent_id: Agent ID to wait for
+            endpoint_name: Endpoint name, defaults to "DEFAULT"
+            max_wait: Maximum wait time in seconds
+
+        Returns:
+            Agent endpoint ARN when ready
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait:
+            try:
+                resp = self.client.get_agent_runtime_endpoint(
+                    agentRuntimeId=agent_id,
+                    endpointName=endpoint_name,
+                )
+                status = resp.get("status", "UNKNOWN")
+
+                if status == "READY":
+                    return resp["agentRuntimeEndpointArn"]
+                elif status in ["CREATE_FAILED", "UPDATE_FAILED"]:
+                    raise Exception(
+                        f"Agent endpoint {status.lower().replace('_', ' ')}: {resp.get('failureReason', 'Unknown')}"
+                    )
+                elif status not in ["CREATING", "UPDATING"]:
+                    pass
+            except self.client.exceptions.ResourceNotFoundException:
+                pass
+            except Exception as e:
+                if "ResourceNotFoundException" not in str(e):
+                    raise
+            time.sleep(2)
+        return (
+            f"Endpoint is taking longer than {max_wait} seconds to be ready, "
+            f"please check status and try to invoke after some time"
+        )
+
+    def get_agent_runtime(self, agent_id: str) -> Dict:
+        """Get agent runtime details.
+
+        Args:
+            agent_id: Agent ID to get details for
+
+        Returns:
+            Agent runtime details
+        """
+        return self.client.get_agent_runtime(agentRuntimeId=agent_id)
+
+    def get_agent_runtime_endpoint(self, agent_id: str, endpoint_name: str = "DEFAULT") -> Dict:
+        """Get agent runtime endpoint details.
+
+        Args:
+            agent_id: Agent ID to get endpoint for
+            endpoint_name: Endpoint name, defaults to "DEFAULT"
+
+        Returns:
+            Agent endpoint details
+        """
+        return self.client.get_agent_runtime_endpoint(
+            agentRuntimeId=agent_id,
+            endpointName=endpoint_name,
+        )
+
+    def invoke_endpoint(
+        self,
+        agent_arn: str,
+        payload: str,
+        session_id: str,
+        endpoint_name: str = "DEFAULT",
+        user_id: Optional[str] = None,
+    ) -> Dict:
+        """Invoke agent endpoint."""
+        req = {
+            "agentRuntimeArn": agent_arn,
+            "qualifier": endpoint_name,
+            "runtimeSessionId": session_id,
+            "payload": payload,
+        }
+
+        if user_id:
+            req["runtimeUserId"] = user_id
+
+        response = self.dataplane_client.invoke_agent_runtime(**req)
+        return _handle_aws_response(response)
+
+
+class HttpBedrockAgentCoreClient:
+    """Bedrock AgentCore client for agent management using HTTP requests with bearer token."""
+
+    def __init__(self, region: str):
+        """Initialize HttpBedrockAgentCoreClient.
+
+        Args:
+            region: AWS region for the client
+        """
+        self.region = region
+        self.dp_endpoint = get_data_plane_endpoint(region)
+        self.logger = logging.getLogger(f"bedrock_agentcore.http_runtime.{region}")
+
+        self.logger.debug("Initializing HTTP Bedrock AgentCore client for region: %s", region)
+        self.logger.debug("Data plane: %s", self.dp_endpoint)
+
+    def invoke_endpoint(
+        self,
+        agent_arn: str,
+        payload,
+        session_id: str,
+        bearer_token: Optional[str],
+        endpoint_name: str = "DEFAULT",
+    ) -> Dict:
+        """Invoke agent endpoint using HTTP request with bearer token.
+
+        Args:
+            agent_arn: Agent ARN to invoke
+            payload: Payload to send (dict or string)
+            session_id: Session ID for the request
+            bearer_token: Bearer token for authentication
+            endpoint_name: Endpoint name, defaults to "DEFAULT"
+
+        Returns:
+            Response from the agent endpoint
+        """
+        # Escape agent ARN for URL
+        escaped_arn = urllib.parse.quote(agent_arn, safe="")
+
+        # Build URL
+        url = f"{self.dp_endpoint}/runtimes/{escaped_arn}/invocations"
+        # Headers
+        headers = {
+            "Authorization": f"Bearer {bearer_token}",
+            "Content-Type": "application/json",
+            "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": session_id,
+        }
+
+        # Parse the payload string back to JSON object to send properly
+        # This ensures consistent payload structure between boto3 and HTTP clients
+        try:
+            body = json.loads(payload) if isinstance(payload, str) else payload
+        except json.JSONDecodeError:
+            # Fallback for non-JSON strings - wrap in payload object
+            self.logger.warning("Failed to parse payload as JSON, wrapping in payload object")
+            body = {"payload": payload}
+
+        try:
+            # Make request with timeout
+            response = requests.post(
+                url,
+                params={"qualifier": endpoint_name},
+                headers=headers,
+                json=body,
+                timeout=100,
+                stream=True,
+            )
+            return _handle_http_response(response)
+        except requests.exceptions.RequestException as e:
+            self.logger.error("Failed to invoke agent endpoint: %s", str(e))
+            raise
+
+
+class LocalBedrockAgentCoreClient:
+    """Local Bedrock AgentCore client for invoking endpoints."""
+
+    def __init__(self, endpoint: str):
+        """Initialize the local client with the given endpoint."""
+        self.endpoint = endpoint
+        self.logger = logging.getLogger("bedrock_agentcore.http_local")
+
+    def invoke_endpoint(self, session_id: str, payload: str, workload_access_token: str):
+        """Invoke the endpoint with the given parameters."""
+        from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER
+
+        url = f"{self.endpoint}/invocations"
+
+        headers = {
+            "Content-Type": "application/json",
+            ACCESS_TOKEN_HEADER: workload_access_token,
+            SESSION_HEADER: session_id,
+        }
+
+        try:
+            body = json.loads(payload) if isinstance(payload, str) else payload
+        except json.JSONDecodeError:
+            # Fallback for non-JSON strings - wrap in payload object
+            self.logger.warning("Failed to parse payload as JSON, wrapping in payload object")
+            body = {"payload": payload}
+
+        try:
+            # Make request with timeout
+            response = requests.post(url, headers=headers, json=body, timeout=100, stream=True)
+            return _handle_http_response(response)
+        except requests.exceptions.RequestException as e:
+            self.logger.error("Failed to invoke agent endpoint: %s", str(e))
+            raise

bedrock_agentcore_starter_toolkit/utils/endpoints.py

@@ -0,0 +1,32 @@
+"""Endpoint utilities for BedrockAgentCore services."""
+
+import os
+
+# Environment-configurable constants with fallback defaults
+DP_ENDPOINT_OVERRIDE = os.getenv("BEDROCK_AGENTCORE_DP_ENDPOINT")
+CP_ENDPOINT_OVERRIDE = os.getenv("BEDROCK_AGENTCORE_CP_ENDPOINT")
+DEFAULT_REGION = os.getenv("AWS_REGION", "us-west-2")
+
+
+def get_data_plane_endpoint(region: str = DEFAULT_REGION) -> str:
+    """Get the data plane endpoint URL for BedrockAgentCore services.
+
+    Args:
+        region: AWS region to use. Defaults to DEFAULT_REGION.
+
+    Returns:
+        The data plane endpoint URL, either from environment override or constructed URL.
+    """
+    return DP_ENDPOINT_OVERRIDE or f"https://bedrock-agentcore.{region}.amazonaws.com"
+
+
+def get_control_plane_endpoint(region: str = DEFAULT_REGION) -> str:
+    """Get the control plane endpoint URL for BedrockAgentCore services.
+
+    Args:
+        region: AWS region to use. Defaults to DEFAULT_REGION.
+
+    Returns:
+        The control plane endpoint URL, either from environment override or constructed URL.
+    """
+    return CP_ENDPOINT_OVERRIDE or f"https://bedrock-agentcore-control.{region}.amazonaws.com"

bedrock_agentcore_starter_toolkit/utils/logging_config.py

@@ -0,0 +1,72 @@
+"""Centralized logging configuration for bedrock-agentcore-starter-toolkit."""
+
+import logging
+
+_LOGGING_CONFIGURED = False
+
+
+def setup_toolkit_logging(mode: str = "sdk") -> None:
+    """Setup logging for bedrock-agentcore-starter-toolkit.
+
+    Args:
+        mode: "cli" or "sdk" (defaults to "sdk")
+    """
+    global _LOGGING_CONFIGURED
+    if _LOGGING_CONFIGURED:
+        return  # Already configured, prevent duplicates
+
+    if mode == "cli":
+        _setup_cli_logging()
+    elif mode == "sdk":
+        _setup_sdk_logging()
+    else:
+        raise ValueError(f"Invalid logging mode: {mode}. Must be 'cli' or 'sdk'")
+
+    _LOGGING_CONFIGURED = True
+
+
+def _setup_cli_logging() -> None:
+    """Setup logging for CLI usage with RichHandler."""
+    try:
+        from rich.logging import RichHandler
+
+        from ..cli.common import console
+
+        FORMAT = "%(message)s"
+        logging.basicConfig(
+            level="INFO",
+            format=FORMAT,
+            handlers=[RichHandler(show_time=False, show_path=False, show_level=False, console=console)],
+            force=True,  # Override any existing configuration
+        )
+    except ImportError:
+        # Fallback if rich is not available
+        _setup_basic_logging()
+
+
+def _setup_sdk_logging() -> None:
+    """Setup logging for SDK usage (notebooks, scripts, imports) with StreamHandler."""
+    # Configure logger for ALL toolkit modules (ensures all operation logs appear)
+    toolkit_logger = logging.getLogger("bedrock_agentcore_starter_toolkit")
+
+    if not toolkit_logger.handlers:
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter("%(message)s"))
+        toolkit_logger.addHandler(handler)
+        toolkit_logger.setLevel(logging.INFO)
+
+
+def _setup_basic_logging() -> None:
+    """Setup basic logging as fallback."""
+    logging.basicConfig(level=logging.INFO, format="%(message)s", force=True)
+
+
+def is_logging_configured() -> bool:
+    """Check if toolkit logging has been configured."""
+    return _LOGGING_CONFIGURED
+
+
+def reset_logging_config() -> None:
+    """Reset logging configuration state (for testing)."""
+    global _LOGGING_CONFIGURED
+    _LOGGING_CONFIGURED = False