amazon-bedrock-agentcore-samples/02-use-cases/SRE-agent/sre_agent/constants.py

#!/usr/bin/env python3

import logging
from typing import Optional
from pydantic import BaseModel, Field

# Configure logging with basicConfig
logging.basicConfig(
    level=logging.INFO,  # Set the log level to INFO
    # Define log message format
    format="%(asctime)s,p%(process)s,{%(filename)s:%(lineno)d},%(levelname)s,%(message)s",
)

logger = logging.getLogger(__name__)


class ModelConfig(BaseModel):
    """Model configuration constants."""

    # Anthropic model IDs
    anthropic_model_id: str = Field(
        default="claude-sonnet-4-20250514",
        description="Default Anthropic Claude model ID",
    )

    # Amazon Bedrock model IDs
    bedrock_model_id: str = Field(
        default="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
        description="Default Amazon Bedrock Claude model ID",
    )

    # Model parameters
    default_temperature: float = Field(
        default=0.1,
        ge=0.0,
        le=2.0,
        description="Default temperature for LLM generation",
    )

    default_max_tokens: int = Field(
        default=4096,
        ge=1,
        le=100000,
        description="Default max tokens for agent responses",
    )

    output_formatter_max_tokens: int = Field(
        default=1000,
        ge=1,
        le=100000,
        description="Max tokens for output formatter LLM calls",
    )


class AWSConfig(BaseModel):
    """AWS configuration constants."""

    default_region: str = Field(default="us-east-1", description="Default AWS region")

    bedrock_endpoint_url: str = Field(
        default="https://bedrock-agentcore-control.us-east-1.amazonaws.com",
        description="Amazon Bedrock AgentCore control endpoint URL",
    )

    credential_provider_endpoint_url: str = Field(
        default="https://us-east-1.prod.agent-credential-provider.cognito.aws.dev",
        description="AWS credential provider endpoint URL",
    )


class TimeoutConfig(BaseModel):
    """Timeout configuration constants."""

    graph_execution_timeout_seconds: int = Field(
        default=600,
        ge=1,
        le=3600,
        description="Maximum time to wait for graph execution (10 minutes)",
    )

    mcp_tools_timeout_seconds: int = Field(
        default=30,
        ge=1,
        le=300,
        description="Maximum time to wait for MCP tools loading",
    )


class PromptConfig(BaseModel):
    """Prompt configuration constants."""

    prompts_directory: str = Field(
        default="config/prompts",
        description="Directory containing prompt template files",
    )

    agent_prompt_files: dict[str, str] = Field(
        default={
            "kubernetes": "kubernetes_agent_prompt.txt",
            "logs": "logs_agent_prompt.txt",
            "metrics": "metrics_agent_prompt.txt",
            "runbooks": "runbooks_agent_prompt.txt",
        },
        description="Mapping of agent types to their prompt files",
    )

    supervisor_prompt_files: dict[str, str] = Field(
        default={
            "plan_aggregation": "supervisor_plan_aggregation.txt",
            "standard_aggregation": "supervisor_standard_aggregation.txt",
            "system": "supervisor_aggregation_system.txt",
        },
        description="Supervisor aggregation prompt files",
    )

    output_formatter_prompt_files: dict[str, str] = Field(
        default={
            "executive_summary_system": "executive_summary_system.txt",
            "executive_summary_user_template": "executive_summary_user_template.txt",
        },
        description="Output formatter prompt files",
    )

    base_prompt_file: str = Field(
        default="agent_base_prompt.txt",
        description="Base prompt template used by all agents",
    )

    enable_prompt_caching: bool = Field(
        default=True, description="Whether to enable LRU caching for prompt loading"
    )

    max_cache_size: int = Field(
        default=32,
        ge=1,
        le=128,
        description="Maximum number of prompts to cache in memory",
    )


class ApplicationConfig(BaseModel):
    """Application configuration constants."""

    agent_model_name: str = Field(
        default="sre-multi-agent", description="Model name returned in API responses"
    )

    default_output_dir: str = Field(
        default="./reports",
        description="Default directory for saving investigation reports",
    )

    conversation_state_file: str = Field(
        default=".multi_agent_conversation_state.json",
        description="Filename for saving conversation state",
    )

    spinner_chars: list[str] = Field(
        default=["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"],
        description="Characters used for spinner animation",
    )


class SREConstants:
    """Central constants configuration for the SRE Agent system.

    This class provides a centralized way to access all configuration constants
    used throughout the SRE Agent application. It uses Pydantic models for
    validation and type safety.

    Usage:
        from .constants import SREConstants

        # Access model configuration
        model_id = SREConstants.model.anthropic_model_id
        temperature = SREConstants.model.default_temperature

        # Access AWS configuration
        region = SREConstants.aws.default_region

        # Access timeout configuration
        timeout = SREConstants.timeouts.graph_execution_timeout_seconds

        # Access prompt configuration
        prompts_dir = SREConstants.prompts.prompts_directory
        agent_files = SREConstants.prompts.agent_prompt_files

        # Access application configuration
        output_dir = SREConstants.app.default_output_dir
    """

    model: ModelConfig = ModelConfig()
    aws: AWSConfig = AWSConfig()
    timeouts: TimeoutConfig = TimeoutConfig()
    prompts: PromptConfig = PromptConfig()
    app: ApplicationConfig = ApplicationConfig()

    @classmethod
    def get_model_config(cls, provider: str, **kwargs) -> dict:
        """Get model configuration for a specific provider.

        Args:
            provider: LLM provider ("anthropic" or "bedrock")
            **kwargs: Additional configuration overrides

        Returns:
            Dictionary with model configuration
        """
        if provider == "anthropic":
            return {
                "model_id": kwargs.get("model_id", cls.model.anthropic_model_id),
                "max_tokens": kwargs.get("max_tokens", cls.model.default_max_tokens),
                "temperature": kwargs.get("temperature", cls.model.default_temperature),
            }
        elif provider == "bedrock":
            return {
                "model_id": kwargs.get("model_id", cls.model.bedrock_model_id),
                "region_name": kwargs.get("region_name", cls.aws.default_region),
                "max_tokens": kwargs.get("max_tokens", cls.model.default_max_tokens),
                "temperature": kwargs.get("temperature", cls.model.default_temperature),
            }
        else:
            raise ValueError(f"Unsupported provider: {provider}")

    @classmethod
    def get_output_formatter_config(cls, provider: str, **kwargs) -> dict:
        """Get model configuration for output formatter.

        Args:
            provider: LLM provider ("anthropic" or "bedrock")
            **kwargs: Additional configuration overrides

        Returns:
            Dictionary with output formatter model configuration
        """
        config = cls.get_model_config(provider, **kwargs)
        # Override max_tokens for output formatter
        config["max_tokens"] = kwargs.get(
            "max_tokens", cls.model.output_formatter_max_tokens
        )
        return config

    @classmethod
    def get_prompt_config(cls) -> PromptConfig:
        """Get prompt configuration.

        Returns:
            PromptConfig instance with all prompt settings
        """
        return cls.prompts


# Convenience instance for easy access
constants = SREConstants()

# Legacy support - individual constants for backward compatibility if needed
ANTHROPIC_MODEL_ID = constants.model.anthropic_model_id
BEDROCK_MODEL_ID = constants.model.bedrock_model_id
DEFAULT_TEMPERATURE = constants.model.default_temperature
DEFAULT_MAX_TOKENS = constants.model.default_max_tokens
DEFAULT_AWS_REGION = constants.aws.default_region
GRAPH_EXECUTION_TIMEOUT_SECONDS = constants.timeouts.graph_execution_timeout_seconds
AGENT_MODEL_NAME = constants.app.agent_model_name
DEFAULT_OUTPUT_DIR = constants.app.default_output_dir