utils.py

"""
This file contains utility functions used across the codebase.
"""
import json
import yaml
import logging
from constants import *
from pathlib import Path
from typing import Any, Dict, Optional

# Global config cache
_config: Optional[Dict[str, Any]] = None

# Global pricing data cache
_pricing_data: Optional[Dict[str, Any]] = None

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s,p%(process)s,{%(filename)s:%(lineno)d},%(levelname)s,%(message)s",
)
logger = logging.getLogger(__name__)


def load_config() -> Dict[str, Any]:
    """Load configuration from config.yaml file.

    Uses CONFIG_FILE_NAME from constants.py to locate the config file.
    Caches the configuration for subsequent calls.

    Returns:
        Dictionary with configuration data

    Raises:
        FileNotFoundError: If config file doesn't exist
        yaml.YAMLError: If config file is malformed
    """
    global _config

    if _config is not None:
        return _config

    config_file = Path(CONFIG_FILE_NAME)
    if not config_file.exists():
        error_msg = f"Config file not found: {CONFIG_FILE_NAME}"
        logger.error(error_msg)
        raise FileNotFoundError(error_msg)

    try:
        with open(config_file, "r", encoding="utf-8") as f:
            _config = yaml.safe_load(f)
        logger.debug(f"Loaded configuration from {CONFIG_FILE_NAME}")
        return _config
    except yaml.YAMLError as e:
        logger.error(f"Failed to parse config file: {e}", exc_info=True)
        raise
    except Exception as e:
        logger.error(f"Failed to load config file: {e}", exc_info=True)
        raise

def load_pricing_data() -> Dict[str, Any]:
    """Load pricing data from YAML file.

    Returns:
        Dictionary with pricing data
    """
    global _pricing_data
    if _pricing_data is not None:
        return _pricing_data

    # Get pricing file path from config
    config = load_config()
    pricing_fpath = config.get("file_paths", {}).get("model_pricing", "model_pricing.yaml")
    pricing_file = Path(pricing_fpath)

    if not pricing_file.exists():
        logger.warning("Pricing file not found, using default values")
        raise FileNotFoundError(f"Pricing file not found: {pricing_fpath}")

    try:
        with open(pricing_file, "r", encoding="utf-8") as f:
            _pricing_data = yaml.safe_load(f)
        logger.debug(f"Loaded pricing data from {pricing_fpath}")
        return _pricing_data
    except Exception as e:
        logger.error(f"Failed to load pricing data: {e}", exc_info=True)
        return {}


def calculate_cost(
    model_id: str,
    input_tokens: int,
    output_tokens: int,
) -> float:
    """Calculate estimated cost based on token usage.

    Args:
        model_id: Amazon Bedrock model ID
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        cache_read_tokens: Number of cache read tokens (prompt caching)
        cache_write_tokens: Number of cache write tokens (prompt caching)

    Returns:
        Estimated cost in dollars
    """
    pricing_data = load_pricing_data()
    # Map Bedrock model ID to pricing model name
    model_mapping = pricing_data.get("bedrock_model_mapping", {})
    pricing_model = model_mapping.get(model_id)
    if not pricing_model:
        logger.error(f"No pricing mapping for model {model_id}, provide information about the model pricing in the pricing file first.")
        raise
    # Get pricing for the model
    models_pricing = pricing_data.get("models", {})
    model_pricing = models_pricing.get(pricing_model)
    # Calculate costs (pricing is per 1000 tokens)
    input_cost = (input_tokens / 1000) * model_pricing.get("input")
    output_cost = (output_tokens / 1000) * model_pricing.get("output")
    total_cost = input_cost + output_cost
    return total_cost


def call_tool(
    tool_name: str,
    arguments: Dict[str, Any],
) -> Dict[str, Any]:
    """Call a tool on the Code Interpreter using invoke_code_interpreter API.

    This is the core helper function that wraps the boto3 invoke_code_interpreter call.
    It follows the pattern from the AWS AgentCore samples.

    Args:
        tool_name: Name of the tool to call (e.g., 'executeCommand', 'writeFiles', 'executeCode')
        arguments: Dictionary of arguments to pass to the tool

    Returns:
        Dictionary with the tool result

    Example:
        result = call_tool("executeCommand", {"command": "pip install requests"})
    """
    # Import globals from code_exec_with_mcp_agent module
    import code_exec_with_mcp_agent

    _dp_client = code_exec_with_mcp_agent._dp_client
    _interpreter_id = code_exec_with_mcp_agent._interpreter_id
    _session_id = code_exec_with_mcp_agent._session_id

    if not _dp_client or not _interpreter_id or not _session_id:
        raise RuntimeError(
            "Code Interpreter not initialized. Call _get_or_create_code_client first."
        )
    logger.debug(f"Calling tool: {tool_name}")
    logger.debug(f"Arguments: {json.dumps(arguments, indent=2, default=str)}")
    try:
        # In this case, we are calling or executing the <tool_name> using the same
        # interpreter and session ID to maintain session persistence.
        response = _dp_client.invoke_code_interpreter(
            codeInterpreterIdentifier=_interpreter_id,
            sessionId=_session_id,
            name=tool_name,
            arguments=arguments,
        )
        # Process the stream and extract result
        for event in response["stream"]:
            if "result" in event:
                result = event["result"]
                logger.debug(f"Tool result: {json.dumps(result, indent=2, default=str)}")
                return result
        # If no result found in stream
        logger.warning(f"No result found in stream for tool: {tool_name}")
        return {"isError": True, "content": [{"text": "No result in response stream"}]}
    except Exception as e:
        logger.error(f"Error calling tool {tool_name}: {e}", exc_info=True)
        return {"isError": True, "content": [{"text": f"Tool execution error: {str(e)}"}]}


def collect_mcp_registry_files() -> list[dict[str, str]]:
    """
    Collect all MCP registry files for upload to sandbox.

    Returns:
        List of dicts with 'path' and 'text' keys for each file
    """
    logger.info("Collecting MCP registry files for upload")

    # Get MCP registry path from config
    config = load_config()
    mcp_registry_path = config.get("mcp_registry", {}).get("registry_path", "mcp_registry")

    files_to_upload = []
    mcp_path = Path(mcp_registry_path)

    if not mcp_path.exists():
        logger.warning(f"MCP registry path not found: {mcp_registry_path}")
        return []

    # Walk through all Python files in mcp_registry
    for py_file in mcp_path.rglob("*.py"):
        try:
            with open(py_file, "r", encoding="utf-8") as f:
                content = f.read()

            # Create relative path for sandbox
            # py_file is already relative to cwd, so just use it as-is
            relative_path = py_file

            files_to_upload.append(
                {
                    "path": str(relative_path),
                    "text": content,
                }
            )
            logger.debug(f"Added file: {relative_path}")

        except Exception as e:
            logger.warning(f"Failed to read {py_file}: {e}")
            continue

    logger.info(f"Collected {len(files_to_upload)} files for upload")
    return files_to_upload

def build_system_prompt() -> str:
    """Build system prompt from file.

    Loads the system prompt from the path defined in config.yaml.

    Returns:
        System prompt string
    """
    # Get system prompt path from config
    config = load_config()
    prompt_path_str = config.get("file_paths", {}).get("system_prompt", "system_prompts/code_exec_prompt.txt")

    logger.info(f"Loading system prompt from {prompt_path_str}")

    prompt_path = Path(prompt_path_str)
    if not prompt_path.exists():
        error_msg = f"System prompt file not found at {prompt_path_str}"
        logger.error(error_msg)
        raise FileNotFoundError(error_msg)

    try:
        with open(prompt_path, "r", encoding="utf-8") as f:
            system_prompt = f.read()

        logger.debug(f"Loaded system prompt ({len(system_prompt)} chars)")
        return system_prompt

    except Exception as e:
        logger.error(f"Failed to load system prompt: {e}", exc_info=True)
        raise