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

bedrock_agentcore_starter_toolkit/operations/runtime/models.py

@@ -0,0 +1,76 @@
+"""Pydantic models for operation requests and responses."""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from ...utils.runtime.container import ContainerRuntime
+
+
+# Configure operation models
+class ConfigureResult(BaseModel):
+    """Result of configure operation."""
+
+    config_path: Path = Field(..., description="Path to configuration file")
+    dockerfile_path: Path = Field(..., description="Path to generated Dockerfile")
+    dockerignore_path: Optional[Path] = Field(None, description="Path to generated .dockerignore")
+    runtime: str = Field(..., description="Container runtime name")
+    region: str = Field(..., description="AWS region")
+    account_id: str = Field(..., description="AWS account ID")
+    execution_role: Optional[str] = Field(None, description="AWS execution role ARN")
+    ecr_repository: Optional[str] = Field(None, description="ECR repository URI")
+    auto_create_ecr: bool = Field(False, description="Whether ECR will be auto-created")
+
+
+# Launch operation models
+class LaunchResult(BaseModel):
+    """Result of launch operation."""
+
+    mode: str = Field(..., description="Launch mode: local, push-ecr, or cloud")
+    tag: str = Field(..., description="Docker image tag")
+    env_vars: Optional[Dict[str, str]] = Field(default=None, description="Environment variables for local deployment")
+
+    # Local mode fields
+    port: Optional[int] = Field(default=None, description="Port for local deployment")
+    runtime: Optional[ContainerRuntime] = Field(default=None, description="Container runtime instance")
+
+    # Cloud mode fields
+    ecr_uri: Optional[str] = Field(default=None, description="ECR repository URI")
+    agent_id: Optional[str] = Field(default=None, description="BedrockAgentCore agent ID")
+    agent_arn: Optional[str] = Field(default=None, description="BedrockAgentCore agent ARN")
+
+    # Build output (optional)
+    build_output: Optional[List[str]] = Field(default=None, description="Docker build output")
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)  # For runtime field
+
+
+class InvokeResult(BaseModel):
+    """Result of invoke operation."""
+
+    response: Dict[str, Any] = Field(..., description="Response from Bedrock AgentCore endpoint")
+    session_id: str = Field(..., description="Session ID used for invocation")
+    agent_arn: Optional[str] = Field(default=None, description="BedrockAgentCore agent ARN")
+
+
+# Status operation models
+class StatusConfigInfo(BaseModel):
+    """Configuration information for status."""
+
+    name: str = Field(..., description="Bedrock AgentCore application name")
+    entrypoint: str = Field(..., description="Entrypoint file path")
+    region: Optional[str] = Field(None, description="AWS region")
+    account: Optional[str] = Field(None, description="AWS account ID")
+    execution_role: Optional[str] = Field(None, description="AWS execution role ARN")
+    ecr_repository: Optional[str] = Field(None, description="ECR repository URI")
+    agent_id: Optional[str] = Field(None, description="BedrockAgentCore agent ID")
+    agent_arn: Optional[str] = Field(None, description="BedrockAgentCore agent ARN")
+
+
+class StatusResult(BaseModel):
+    """Result of status operation."""
+
+    config: StatusConfigInfo = Field(..., description="Configuration information")
+    agent: Optional[Dict[str, Any]] = Field(None, description="Agent runtime details or error")
+    endpoint: Optional[Dict[str, Any]] = Field(None, description="Endpoint details or error")

bedrock_agentcore_starter_toolkit/operations/runtime/status.py

@@ -0,0 +1,66 @@
+"""Status operations for Bedrock AgentCore SDK."""
+
+from pathlib import Path
+from typing import Optional
+
+from ...services.runtime import BedrockAgentCoreClient
+from ...utils.runtime.config import load_config
+from .models import StatusConfigInfo, StatusResult
+
+
+def get_status(config_path: Path, agent_name: Optional[str] = None) -> StatusResult:
+    """Get Bedrock AgentCore status including config and runtime details.
+
+    Args:
+        config_path: Path to BedrockAgentCore configuration file
+        agent_name: Name of agent to get status for (for project configurations)
+
+    Returns:
+        StatusResult with config, agent, and endpoint status
+
+    Raises:
+        FileNotFoundError: If configuration file doesn't exist
+        ValueError: If Bedrock AgentCore is not deployed or configuration is invalid
+    """
+    # Load project configuration
+    project_config = load_config(config_path)
+    agent_config = project_config.get_agent_config(agent_name)
+
+    # Build config info
+    config_info = StatusConfigInfo(
+        name=agent_config.name,
+        entrypoint=agent_config.entrypoint,
+        region=agent_config.aws.region,
+        account=agent_config.aws.account,
+        execution_role=agent_config.aws.execution_role,
+        ecr_repository=agent_config.aws.ecr_repository,
+        agent_id=agent_config.bedrock_agentcore.agent_id,
+        agent_arn=agent_config.bedrock_agentcore.agent_arn,
+    )
+
+    # Initialize status result
+    agent_details = None
+    endpoint_details = None
+
+    # If agent is deployed, get runtime status
+    if agent_config.bedrock_agentcore.agent_id and agent_config.aws.region:
+        try:
+            client = BedrockAgentCoreClient(agent_config.aws.region)
+
+            # Get agent runtime details
+            try:
+                agent_details = client.get_agent_runtime(agent_config.bedrock_agentcore.agent_id)
+            except Exception as e:
+                agent_details = {"error": str(e)}
+
+            # Get endpoint details
+            try:
+                endpoint_details = client.get_agent_runtime_endpoint(agent_config.bedrock_agentcore.agent_id)
+            except Exception as e:
+                endpoint_details = {"error": str(e)}
+
+        except Exception as e:
+            agent_details = {"error": f"Failed to initialize Bedrock AgentCore client: {e}"}
+            endpoint_details = {"error": f"Failed to initialize Bedrock AgentCore client: {e}"}
+
+    return StatusResult(config=config_info, agent=agent_details, endpoint=endpoint_details)

bedrock_agentcore_starter_toolkit/services/ecr.py

@@ -0,0 +1,55 @@
+"""ECR (Elastic Container Registry) service integration."""
+
+import base64
+
+import boto3
+
+from ..utils.runtime.container import ContainerRuntime
+
+
+def get_account_id() -> str:
+    """Get AWS account ID."""
+    return boto3.client("sts").get_caller_identity()["Account"]
+
+
+def get_region() -> str:
+    """Get AWS region."""
+    return boto3.Session().region_name or "us-west-2"
+
+
+def create_ecr_repository(repo_name: str, region: str) -> str:
+    """Create or get existing ECR repository."""
+    ecr = boto3.client("ecr", region_name=region)
+    try:
+        response = ecr.create_repository(repositoryName=repo_name)
+        return response["repository"]["repositoryUri"]
+    except ecr.exceptions.RepositoryAlreadyExistsException:
+        response = ecr.describe_repositories(repositoryNames=[repo_name])
+        return response["repositories"][0]["repositoryUri"]
+
+
+def deploy_to_ecr(local_tag: str, repo_name: str, region: str, container_runtime: ContainerRuntime) -> str:
+    """Build and push image to ECR."""
+    ecr = boto3.client("ecr", region_name=region)
+
+    # Get or create repository
+    ecr_uri = create_ecr_repository(repo_name, region)
+
+    # Get auth token
+    auth_data = ecr.get_authorization_token()["authorizationData"][0]
+    token = base64.b64decode(auth_data["authorizationToken"]).decode("utf-8")
+    username, password = token.split(":")
+
+    # Login to ECR
+    if not container_runtime.login(auth_data["proxyEndpoint"], username, password):
+        raise RuntimeError("Failed to login to ECR")
+
+    # Tag and push
+    ecr_tag = f"{ecr_uri}:latest"
+    if not container_runtime.tag(local_tag, ecr_tag):
+        raise RuntimeError("Failed to tag image")
+
+    if not container_runtime.push(ecr_tag):
+        raise RuntimeError("Failed to push image to ECR")
+
+    return ecr_tag

bedrock_agentcore_starter_toolkit/services/runtime.py

@@ -0,0 +1,383 @@
+"""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 ..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,
+    ) -> 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 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 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,
+    ) -> 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
+        )
+
+    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"