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 typing import Dict, List, Optional
from jinja2 import Template
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from dotenv import load_dotenv

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

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

# 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: Optional[AgentConfig] = 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()

        # 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 Python executor
        self.python_executor = PythonExecutor()

        # 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) -> Optional[str]:
        """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 should_continue(self, state: AgentState) -> str:
        """Decide whether to continue executing or end the workflow."""
        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:
            return "end"

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

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

        # Check if the finish() tool has been called
        if "<solution>" in last_message and "</solution>" in last_message:
            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:
            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) -> Optional[Dict]:
        """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 get_trace(self) -> Dict:
        """Get the complete trace of the last execution."""
        if not self.workflow_engine:
            return {}

        return {
            "execution_time": time.strftime('%Y-%m-%d %H:%M:%S'),
            "config": {
                "max_steps": self.config.max_steps,
                "timeout_seconds": self.config.timeout_seconds,
                "verbose": self.config.verbose
            },
            "messages": self.workflow_engine.message_history,
            "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") -> 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

        Returns:
            The final response content
        """
        # 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 logging_sink import get_log_sink, persist_trace_safe
            run_id = time.strftime("%Y%m%d_%H%M%S")
            persist_trace_safe(get_log_sink(), run_id, self.get_trace())
        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.")