src/agent.py

"""
CodeAgent: A LangGraph-based agent for executing Python code and using tools.
Fully modular version with unified tool management.
"""

import os
import re
import time

from dotenv import load_dotenv
from jinja2 import Template
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

from core.access_control import principal_trace_fields
from core.constants import SYSTEM_PROMPT_TEMPLATE
from core.run_metrics import compute_run_metrics

# Import core types anxd constants
from core.types import AgentConfig, AgentState

# Import managers (organized by subsystem)
from managers import (
    ConsoleDisplay,
    Executor,
    # Support
    PackageManager,
    # Workflow
    PlanManager,
    # Execution
    Timing,
    # Tools
    ToolManager,
    ToolSource,
    WorkflowEngine,
    create_agent_state,
    get_executor,
)

# Load environment variables
load_dotenv("./.env")


def get_system_prompt(
    functions: dict[str, dict],
    packages: dict[str, str] = None,
    datasets: dict[str, dict] = None,
) -> str:
    """Generate system prompt using template and functions."""
    if packages is None:
        from core.constants import LIBRARY_CONTENT_DICT

        packages = LIBRARY_CONTENT_DICT
    if datasets is None:
        from src.datasets.registry import get_registry, is_advertised

        _reg = get_registry()
        datasets = {}
        for _did in _reg.list():
            _raw = _reg.get(_did) or {}
            _expr = _raw.get("expression_source", {})
            _feat = _raw.get("feature_mapping", {})
            datasets[_did] = {
                "dataset_id": _did,
                "title": _raw.get("title", ""),
                "preprocessing": _raw.get("preprocessing", ""),
                "accession": _raw.get("accession", ""),
                "organism": _raw.get("organism", ""),
                "modality": _raw.get("modality", ""),
                "advertised": is_advertised(_raw.get("modality", "")),
                "data_level": _raw.get("data_level", ""),
                "expression_url": _expr.get("url", ""),
                "feature_id_type": _raw.get("feature_id_type", ""),
                "requires_collapse": _feat.get("requires_collapse", False),
                "group_columns": _raw.get("group_columns", []),
                "default_contrasts": _raw.get("default_contrasts", []),
                "survival_columns": _raw.get("survival_columns") or {},
                "refusal_rules": _raw.get("refusal_rules", []),
                "limitations": _raw.get("limitations", []),
                "reporting_rules": _raw.get("reporting_rules", []),
                "dataset_disclaimer": _raw.get("dataset_disclaimer") or "",
                "has_curated_subset": bool(_raw.get("curated_sample_list")),
                "curated_n": len(_raw.get("curated_sample_list") or []),
                "curated_sample_source": _raw.get("curated_sample_source") or "",
            }
    return Template(SYSTEM_PROMPT_TEMPLATE).render(
        functions=functions, packages=packages, datasets=datasets
    )


class CodeAgent:
    """A code-based agent that can execute Python code and use tools to solve tasks."""

    def __init__(self, model: BaseChatModel, config: AgentConfig | None = None):
        """
        Initialize the CodeAgent.

        Args:
            model: The language model to use for generation
            config: Configuration for the agent
        """
        self.model = model
        self.config = config or AgentConfig()

        # Authenticated caller for the CURRENT run, threaded into the ADR-0008
        # audit trace (ADR-0012 decision 3). Set per-run via run(principal=...)
        # or set_principal(); None → recorded as "anonymous" (no identity was
        # forwarded, e.g. the specialist called without a threaded identity).
        self.principal: str | None = None

        # Initialize modular components
        self.package_manager = PackageManager()
        self.console = ConsoleDisplay()
        self.tool_manager = ToolManager(self.console)
        self.workflow_engine = WorkflowEngine(model, self.config, self.console)

        # Initialize the code executor via the config-selected factory
        # (ADR-0007 Phase 0 seam). Default EXECUTOR=in_process returns a
        # PythonExecutor — identical behavior to constructing it directly; the
        # `sandbox` executor (ADR-0007 Phase 1) will slot in behind this same
        # Executor interface with no agent-loop change.
        self.python_executor: Executor = get_executor()

        # Setup workflow
        self._setup_workflow()

    # ====================
    # WORKFLOW SETUP
    # ====================

    def _setup_workflow(self):
        """Setup the LangGraph workflow using WorkflowEngine."""
        self.workflow_engine.setup_workflow(self.generate, self.execute, self.should_continue)

    # ====================
    # WORKFLOW NODES
    # ====================

    def generate(self, state: AgentState) -> AgentState:
        """Generate response using LLM with tool-aware prompt."""

        all_schemas = self.tool_manager.get_tool_schemas(openai_format=True)
        functions_dict = {schema["function"]["name"]: schema for schema in all_schemas}

        all_packages = self.package_manager.get_all_packages()

        from src.datasets.registry import get_registry, is_advertised

        _reg = get_registry()
        datasets = {}
        for _did in _reg.list():
            _raw = _reg.get(_did) or {}
            _expr = _raw.get("expression_source", {})
            _feat = _raw.get("feature_mapping", {})
            datasets[_did] = {
                "dataset_id": _did,
                "title": _raw.get("title", ""),
                "preprocessing": _raw.get("preprocessing", ""),
                "accession": _raw.get("accession", ""),
                "organism": _raw.get("organism", ""),
                "modality": _raw.get("modality", ""),
                "advertised": is_advertised(_raw.get("modality", "")),
                "data_level": _raw.get("data_level", ""),
                "expression_url": _expr.get("url", ""),
                "feature_id_type": _raw.get("feature_id_type", ""),
                "requires_collapse": _feat.get("requires_collapse", False),
                "group_columns": _raw.get("group_columns", []),
                "default_contrasts": _raw.get("default_contrasts", []),
                "survival_columns": _raw.get("survival_columns") or {},
                "refusal_rules": _raw.get("refusal_rules", []),
                "limitations": _raw.get("limitations", []),
                "reporting_rules": _raw.get("reporting_rules", []),
                "dataset_disclaimer": _raw.get("dataset_disclaimer") or "",
                "has_curated_subset": bool(_raw.get("curated_sample_list")),
                "curated_n": len(_raw.get("curated_sample_list") or []),
                "curated_sample_source": _raw.get("curated_sample_source") or "",
            }

        system_prompt = get_system_prompt(functions_dict, all_packages, datasets)

        # Truncate conversation history to prevent context overflow: keep the
        # original task message plus only the most recent `memory_window`
        # messages. Without this, every step re-sends the full history
        # (including all prior code outputs), so cost grows roughly with the
        # square of the step count.
        all_messages = state["messages"]
        original_task = all_messages[:1]
        recent = all_messages[1:][-self.config.memory_window :]
        # original_task is a HumanMessage. If recent also starts with a
        # HumanMessage (an observation), drop it so we don't send two
        # consecutive Human turns. If recent starts with an AIMessage, leave
        # it — Human(query) -> AI(...) is a valid alternating sequence, and
        # dropping it would erase the agent's last response (the model would
        # just repeat itself with no new progress).
        if recent and isinstance(recent[0], HumanMessage):
            recent = recent[1:]

        history = original_task + recent

        # The system prompt (dataset manifests + tool schemas, ~15K tokens) is
        # identical on every step of a run. Mark it for Anthropic prompt
        # caching so repeat steps pay only ~10% of the input-token cost for
        # this block instead of re-billing it in full each time.
        system_message = SystemMessage(
            content=[
                {
                    "type": "text",
                    "text": system_prompt,
                    "cache_control": {"type": "ephemeral"},
                }
            ]
        )
        messages = [system_message] + history

        # Newer Claude models (4.6+) require conversations to end with a user message.
        # If the last message is AIMessage (e.g. when resuming after a step-limit hit,
        # where should_continue fires after generate but before execute), convert it to
        # HumanMessage so the API call doesn't fail with "assistant prefill" error.
        if messages and isinstance(messages[-1], AIMessage):
            messages = messages[:-1] + [HumanMessage(content=messages[-1].content)]

        from core.perf import timed

        step_no = state.get("step_count", 0) + 1
        with timed("generate", f"step{step_no}"):
            response = self.model.invoke(messages)

        # Cut the text after the </execute> tag, while keeping the </execute> tag
        if "</execute>" in response.content:
            response.content = response.content.split("</execute>")[0] + "</execute>"

        # Parse the response
        msg = str(response.content)
        llm_reply = AIMessage(content=msg.strip())

        # Update step count
        new_step_count = state.get("step_count", 0) + 1

        return create_agent_state(
            messages=[llm_reply],
            step_count=new_step_count,
            error_count=state.get("error_count", 0),
            start_time=state.get("start_time", time.time()),
            current_plan=self._extract_current_plan(msg),
        )

    def _extract_current_plan(self, content: str) -> str | None:
        """Extract the current plan from the agent's response."""
        return PlanManager.extract_plan_from_content(content)

    def execute(self, state: AgentState) -> AgentState:
        """Execute code using persistent Python executor."""
        try:
            last_message = state["messages"][-1].content
            execute_match = re.search(r"<execute>(.*?)</execute>", last_message, re.DOTALL)

            if execute_match:
                code = execute_match.group(1).strip()

                # Re-inject tool functions before every step. The agent's code runs in a
                # persistent namespace, and agent code that does e.g.
                # `from tools.dataset_tools import some_tool` to inspect a function's
                # signature will overwrite that name with the raw FastMCP FunctionTool
                # object, permanently breaking the tool for the rest of the run
                # ('FunctionTool' object is not callable). Restoring the working
                # wrappers here heals any such shadowing every step, at near-zero cost.
                self.python_executor.send_functions(self.get_all_tool_functions())

                # Execute regular code in persistent environment (tools already injected)
                from core.perf import timed

                with timed("execute", f"step{state.get('step_count', 0)}"):
                    result = self.python_executor(code)

                # Include both the code and result in the observation
                obs = f"\n<observation>\nCode Output:\n{result}</observation>"
                # Use HumanMessage for observations — the environment responding to the agent.
                # This keeps human/AI turns strictly alternating, which newer Claude models require.
                return create_agent_state(
                    messages=[HumanMessage(content=obs.strip())],
                    step_count=state.get("step_count", 0),
                    error_count=state.get("error_count", 0),
                    start_time=state.get("start_time", time.time()),
                    current_plan=state.get("current_plan"),
                )
            else:
                return create_agent_state(
                    messages=[HumanMessage(content="<error>No executable code found</error>")],
                    step_count=state.get("step_count", 0),
                    error_count=state.get("error_count", 0) + 1,
                    start_time=state.get("start_time", time.time()),
                    current_plan=state.get("current_plan"),
                )
        except Exception as e:
            return create_agent_state(
                messages=[HumanMessage(content=f"<error>Execution error: {str(e)}</error>")],
                step_count=state.get("step_count", 0),
                error_count=state.get("error_count", 0) + 1,
                start_time=state.get("start_time", time.time()),
                current_plan=state.get("current_plan"),
            )

    def _record_end_reason(self, reason: str) -> None:
        """Stash why the run ended where the UI reads run outcomes."""
        if self.workflow_engine is not None:
            self.workflow_engine.last_end_reason = reason

    def should_continue(self, state: AgentState) -> str:
        """Decide whether to continue executing or end the workflow.

        Records *why* the run ended in ``last_end_reason`` so the UI can tell the
        three exhaustion cases apart. Continuing only helps the step case: a run
        that ran out of wall clock gets nothing from 15 more steps, and telling
        the user otherwise sends them round the same wall again.
        """
        last_message = state["messages"][-1].content
        step_count = state.get("step_count", 0)
        error_count = state.get("error_count", 0)
        start_time = state.get("start_time", time.time())

        # Check for timeout
        if time.time() - start_time > self.config.timeout_seconds:
            self._record_end_reason("timeout")
            return "end"

        # Check for maximum steps
        if step_count >= self.config.max_steps:
            self._record_end_reason("step_limit")
            return "end"

        # Check for too many errors
        if error_count >= self.config.retry_attempts:
            self._record_end_reason("error_limit")
            return "end"

        # Check if the finish() tool has been called
        if "<solution>" in last_message and "</solution>" in last_message:
            self._record_end_reason("solution")
            return "end"

        # Check if there's an execute tag in the last message
        elif "<execute>" in last_message and "</execute>" in last_message:
            return "execute"

        else:
            # No solution, no code to run — the model simply stopped emitting.
            self._record_end_reason("stalled")
            return "end"

    # ====================
    # PACKAGE MANAGEMENT - Delegated to PackageManager
    # ====================

    def add_packages(self, packages: dict[str, str]) -> bool:
        """Add new packages to the available packages."""
        return self.package_manager.add_packages(packages)

    def get_all_packages(self) -> dict[str, str]:
        """Get all available packages (default + custom)."""
        return self.package_manager.get_all_packages()

    # ====================
    # UNIFIED TOOL MANAGEMENT - Delegated to ToolManager
    # ====================

    def add_tool(self, function: callable, name: str = None, description: str = None) -> bool:
        """Add a tool function to the manager."""
        return self.tool_manager.add_tool(function, name, description, ToolSource.LOCAL)

    def remove_tool(self, name: str) -> bool:
        """Remove a tool by name."""
        return self.tool_manager.remove_tool(name)

    def list_tools(self, source: str = "all", include_details: bool = False) -> list[dict]:
        """List all available tools with optional filtering."""
        source_enum = ToolSource.ALL
        if source.lower() in ["local", "decorated", "mcp"]:
            source_enum = ToolSource(source.lower())

        return self.tool_manager.list_tools(source_enum, include_details)

    def search_tools(self, query: str) -> list[dict]:
        """Search tools by name and description."""
        return self.tool_manager.search_tools(query)

    def get_tool_info(self, name: str) -> dict | None:
        """Get detailed information about a specific tool."""
        tool_info = self.tool_manager.get_tool(name)
        if tool_info:
            return {
                "name": tool_info.name,
                "description": tool_info.description,
                "source": tool_info.source.value,
                "server": tool_info.server,
                "module": tool_info.module,
                "has_function": tool_info.function is not None,
                "required_parameters": tool_info.required_parameters,
                "optional_parameters": tool_info.optional_parameters,
            }
        return None

    def get_all_tool_functions(self) -> dict[str, callable]:
        """Get all tool functions as a dictionary."""
        return self.tool_manager.get_all_functions()

    def _prepare_executor(self) -> None:
        """Inject tools and packages into the Python executor before any run."""
        self.python_executor.send_functions(self.get_all_tool_functions())
        self.package_manager.import_packages(self.python_executor)

    # ====================
    # MCP METHODS - Now delegated to ToolManager
    # ====================

    def add_mcp(self, config_path: str = "./mcp_config.yaml") -> None:
        """Add MCP tools from configuration file."""
        self.tool_manager.add_mcp_server(config_path)

    def add_mcp_http(self, url: str, server_name: str = "decouplerpy") -> None:
        """Add MCP tools from a persistent HTTP MCP server (no per-call spawn)."""
        self.tool_manager.add_mcp_http_server(url, server_name=server_name)

    def list_mcp_tools(self) -> list[dict]:
        """List all loaded MCP tools."""
        return self.tool_manager.list_tools(ToolSource.MCP)

    def list_mcp_servers(self) -> dict[str, list[str]]:
        """List all MCP servers and their tools."""
        return self.tool_manager.list_mcp_servers()

    def show_mcp_status(self) -> None:
        """Display detailed MCP status information to the user."""
        self.tool_manager.show_mcp_status()

    def get_mcp_summary(self) -> dict[str, any]:
        """Get a summary of MCP tools for programmatic access."""
        return self.tool_manager.get_mcp_summary()

    # ====================
    # ENHANCED TOOL FEATURES
    # ====================

    def get_tool_statistics(self) -> dict[str, any]:
        """Get comprehensive tool statistics."""
        return self.tool_manager.get_tool_statistics()

    def validate_tools(self) -> dict[str, list[str]]:
        """Validate all tools and return any issues."""
        return self.tool_manager.validate_tools()

    # ====================
    # TOOL SELECTION MANAGEMENT
    # ====================

    def reset_tool_selection(self):
        """Reset the cached tool selection to allow re-selection on next query."""
        self._selected_tools_cache = None
        if self.use_tool_selection:
            self.console.console.print(
                "🔄 Tool selection cache cleared - will re-select tools on next query"
            )

    def get_selected_tools(self):
        """Get the currently selected tools (if any)."""
        return list(self._selected_tools_cache.keys()) if self._selected_tools_cache else None

    # ====================
    # TRACE AND SUMMARY METHODS
    # ====================

    def set_principal(self, principal: str | None) -> None:
        """Set the authenticated caller recorded in the next trace (ADR-0012).

        The orchestrator (front door) establishes identity via HF OAuth and
        forwards it here; the specialist trusts that forwarded principal (it is
        reached only through the orchestrator's service token — ADR-0012
        decision 2). ``None`` clears it back to anonymous.
        """
        self.principal = principal

    def get_trace(self) -> dict:
        """Get the complete trace of the last execution."""
        if not self.workflow_engine:
            return {}

        messages = self.workflow_engine.message_history

        # The Gradio path drives `graph.stream()` directly rather than
        # `run_workflow()`, so it never fills `message_history` — every trace
        # persisted from the UI (i.e. every production run) went to the sink
        # with `"messages": []`. Fall back to the final graph state, which the
        # streaming path does stash, so the persisted trace actually carries the
        # conversation and `n_post_analysis_reads` is computable in prod.
        if not messages:
            messages = [
                self.workflow_engine._serialize_message(m)
                for m in getattr(self.workflow_engine, "last_state_messages", []) or []
            ]

        # Read-only behavioural instrumentation (see src/core/run_metrics.py).
        # Derived from the message history AFTER the run has finished — it
        # observes, never steers. A failure here must not cost us the trace.
        try:
            metrics = compute_run_metrics(messages)
        except Exception as exc:  # pragma: no cover - defensive
            metrics = {"error": f"run_metrics failed: {exc}"}

        return {
            "execution_time": time.strftime("%Y-%m-%d %H:%M:%S"),
            "metrics": {
                "step_count": getattr(self.workflow_engine, "last_step_count", None),
                "end_reason": getattr(self.workflow_engine, "last_end_reason", None),
                **metrics,
            },
            "config": {
                "max_steps": self.config.max_steps,
                "timeout_seconds": self.config.timeout_seconds,
                "verbose": self.config.verbose,
                # ADR-0012 decision 3: attach the authenticated principal + role
                # so the audit trace records WHO ran each analysis (the missing
                # link between the ADR-0008 app trace and identity).
                **principal_trace_fields(self.principal),
            },
            "messages": messages,
            "trace_logs": self.workflow_engine.trace_logs,
        }

    def get_summary(self) -> dict:
        """Get a summary of the last execution."""
        if not self.workflow_engine:
            return {}
        return self.workflow_engine.generate_summary()

    def save_trace(self, filepath: str = None) -> str:
        """Save the trace of the last execution to a file."""
        if not self.workflow_engine:
            raise RuntimeError("No workflow engine available")
        return self.workflow_engine.save_trace_to_file(filepath)

    def save_summary(self, filepath: str = None) -> str:
        """Save the summary of the last execution to a file."""
        if not self.workflow_engine:
            raise RuntimeError("No workflow engine available")
        return self.workflow_engine.save_summary_to_file(filepath)

    # ====================
    # PUBLIC INTERFACE
    # ====================

    def run(
        self,
        query: str,
        save_trace: bool = False,
        save_summary: bool = False,
        trace_dir: str = "traces",
        principal: str | None = None,
    ) -> str:
        """
        Run the agent with a given query using modular components.

        Args:
            query: The task/question to solve
            save_trace: Whether to save the complete trace to a file
            save_summary: Whether to save the execution summary to a file
            trace_dir: Directory to save trace and summary files
            principal: Authenticated caller for this run (ADR-0012); recorded in
                the audit trace. Back-compat default None → "anonymous".

        Returns:
            The final response content
        """
        # Record the authenticated caller for this run so it lands in the trace
        # (ADR-0012). Only overwrite when explicitly passed, so a caller that set
        # the principal via set_principal() before run() is preserved.
        if principal is not None:
            self.principal = principal

        # Start timing the overall execution
        overall_timing = Timing(start_time=time.time())

        # Display task header
        self.console.print_task_header(query)

        # Prepare executor (tools + packages) — shared with Gradio path
        self._prepare_executor()

        # Display tool statistics
        stats = self.tool_manager.get_tool_statistics()
        mcp_servers = self.tool_manager.list_mcp_servers()
        self.console.console.print(f"🛠️  Loaded {stats['total_tools']} total tools:")
        if stats["by_source"]["decorated"] > 0:
            self.console.console.print(f"   🎯 Decorated tools: {stats['by_source']['decorated']}")
        if stats["by_source"]["mcp"] > 0:
            self.console.console.print(
                f"   🔗 MCP tools: {stats['by_source']['mcp']} from {len(mcp_servers)} servers"
            )
            for server_name, tools in mcp_servers.items():
                self.console.console.print(f"      • {server_name}: {len(tools)} tools")

        # Create initial state
        input_state = create_agent_state(
            messages=[HumanMessage(content=query)],
            step_count=0,
            error_count=0,
            start_time=time.time(),
            current_plan=None,
        )

        # Execute workflow using WorkflowEngine and get result with final state
        result, final_state = self.workflow_engine.run_workflow(input_state)

        # Complete overall timing and display summary
        overall_timing.end_time = time.time()

        # Extract final state information for summary
        final_step_count = final_state.get("step_count", 0) if final_state else 0
        final_error_count = final_state.get("error_count", 0) if final_state else 0

        self.console.print_execution_summary(
            final_step_count, final_error_count, overall_timing.duration
        )

        # ALWAYS-ON audit trace persistence via the configured sink (local |
        # hf | s3, env LOG_SINK; default local). This is independent of the
        # opt-in `save_trace` file dump below — on the live path the trace is
        # always persisted so it isn't merely held in memory. The sink write is
        # wrapped so a logging failure never crashes the run.
        try:
            from core.trace_redaction import redact_trace_safe
            from logging_sink import get_log_sink, persist_trace_safe

            run_id = time.strftime("%Y%m%d_%H%M%S")
            # ADR-0013: scrub secrets/credentials/PII on a COPY before it reaches
            # any sink — fail-closed to a minimal trace if redaction itself fails,
            # so an unredacted payload is never persisted. The in-memory trace the
            # UI/eval harness reads is untouched.
            redacted = redact_trace_safe(self.get_trace(), run_id)
            persist_trace_safe(get_log_sink(), run_id, redacted)
        except Exception as e:  # noqa: BLE001 — logging must never crash a run
            print(f"[log_sink] trace persistence skipped: {e}")

        # Save trace and summary to a local file if explicitly requested
        # (backward-compatible opt-in, separate from the always-on sink above)
        if save_trace or save_summary:
            # Create trace directory if it doesn't exist
            from pathlib import Path

            trace_path = Path(trace_dir)
            trace_path.mkdir(parents=True, exist_ok=True)

            if save_trace:
                trace_file = trace_path / f"agent_trace_{time.strftime('%Y%m%d_%H%M%S')}.json"
                saved_trace = self.workflow_engine.save_trace_to_file(str(trace_file))
                self.console.console.print(f"💾 Trace saved to: {saved_trace}")

            if save_summary:
                summary_file = trace_path / f"agent_summary_{time.strftime('%Y%m%d_%H%M%S')}.json"
                saved_summary = self.workflow_engine.save_summary_to_file(str(summary_file))
                self.console.console.print(f"📊 Summary saved to: {saved_summary}")

        return result


# ====================
# EXAMPLE USAGE
# ====================

if __name__ == "__main__":
    """
    CLI entry point for local testing.

    Usage:
        python src/agent.py
        python src/agent.py "Tell me about the Moffitt PDAC dataset."
        python src/agent.py "Run differential expression on the Moffitt dataset."

    Requires:
        - Anthropic_API_KEY environment variable (or .env file at project root)
        - mcp_config.yaml at project root
        - R + limma installed locally for method=limma (optional;
          method=ttest works without R)
    """
    import sys

    from dotenv import load_dotenv
    from langchain_anthropic import ChatAnthropic

    # Load .env from project root (one level up from src/)
    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    load_dotenv(os.path.join(project_root, ".env"))

    # Query: first CLI argument, or a sensible default
    query = (
        " ".join(sys.argv[1:])
        if len(sys.argv) > 1
        else ("Tell me about the Moffitt PDAC dataset and what analyses are supported.")
    )

    api_key = os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("Anthropic_API_KEY")
    if not api_key:
        print("ANTHROPIC_API_KEY not set. Export it or add it to .env at the project root.")
        sys.exit(1)

    model = ChatAnthropic(
        model="claude-sonnet-4-5-20250929",
        temperature=0,
        api_key=api_key,
    )

    config = AgentConfig(
        max_steps=15,
        retry_attempts=3,
        timeout_seconds=1200,
        verbose=True,
    )

    agent = CodeAgent(model=model, config=config)

    # Load MCP tools — config is at project root, not inside src/
    mcp_config = os.path.join(project_root, "mcp_config.yaml")
    if os.path.exists(mcp_config):
        try:
            agent.add_mcp(mcp_config)
            stats = agent.get_tool_statistics()
            print(f"Loaded {stats['total_tools']} tools ({stats['by_source']['mcp']} MCP)")
        except Exception as e:
            print(f"MCP tools could not be loaded: {e}")
            print("   Analysis will run without decoupleRpy tools.")
    else:
        print(f"mcp_config.yaml not found at {mcp_config}")

    print("\nQuery: " + query + "\n" + "-" * 60)
    agent.run(query, save_trace=False, save_summary=False)
    print("-" * 60 + "\nDone.")