"""
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 tag, while keeping the tag
if "" in response.content:
response.content = response.content.split("")[0] + ""
# 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"(.*?)", 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\nCode Output:\n{result}"
# 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="No executable code found")],
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"Execution error: {str(e)}")],
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 "" in last_message and "" in last_message:
self._record_end_reason("solution")
return "end"
# Check if there's an execute tag in the last message
elif "" in last_message and "" 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.")