Upload 56 files

AlphaGenome_mcp.py

@@ -0,0 +1,74 @@
+"""
+Model Context Protocol (MCP) for AlphaGenome
+
+AlphaGenome provides state-of-the-art AI-driven tools for genomic sequence analysis, variant effect prediction, and functional genomics visualization. The platform enables researchers to predict regulatory elements, assess variant impacts, and explore molecular mechanisms across diverse cell types and tissues.
+
+This MCP Server contains the tools extracted from the following tutorials with their tools:
+1. batch_variant_scoring
+    - score_variants_batch: Score multiple genetic variants in batch using configurable variant scorers
+    - filter_variant_scores: Filter variant scores by ontology criteria (e.g., cell types)
+2. essential_commands
+    - create_genomic_interval: Create genomic intervals for DNA regions
+    - create_genomic_variant: Create genomic variants for genetic changes  
+    - create_track_data: Create TrackData objects from arrays and metadata
+    - create_variant_scores: Create AnnData objects for variant scoring results
+    - genomic_interval_operations: Perform operations on genomic intervals
+    - variant_interval_operations: Check variant overlaps with intervals  
+    - track_data_operations: Filter, resize, slice and transform TrackData
+    - track_data_resolution_conversion: Convert between track data resolutions
+3. example_analysis_workflow
+    # REMOVED - Too TAL1-specific:
+    # - visualize_tal1_variant_positions: Visualize genomic positions of oncogenic TAL1 variants
+    # - predict_variant_functional_impact: Predict functional effects of specific TAL1 variants
+    # - compare_oncogenic_vs_background_variants: Compare predicted effects of disease variants vs background
+4. quick_start
+    - predict_dna_sequence: Predict genomic tracks from DNA sequence
+    - predict_genome_interval: Predict genomic tracks for reference genome intervals  
+    # REMOVED - Redundant with visualize_variant_effects:
+    # - predict_variant_effects: Predict and visualize genetic variant effects
+    # REMOVED - Redundant with score_variants_batch:
+    # - score_variant_effect: Score genetic variant effects using variant scorers
+    - ism_analysis: Perform in silico mutagenesis analysis with sequence logos
+    # REMOVED - Should be parameter in other functions:
+    # - mouse_predictions: Make predictions for mouse sequences and intervals
+5. tissue_ontology_mapping
+    - explore_output_metadata: Explore and filter output metadata for specific organisms and search terms
+    - count_tracks_by_output_type: Count tracks by output type for human and mouse organisms
+6. variant_scoring_ui
+    # REMOVED - Redundant with score_variants_batch:
+    # - score_variant: Score a single variant with multiple variant scorers and save results
+    - visualize_variant_effects: Generate comprehensive variant effect visualization across multiple modalities
+7. visualization_modality_tour
+    - visualize_gene_expression: Visualize RNA_SEQ and CAGE gene expression predictions
+    # REMOVED - Redundant with visualize_variant_effects:
+    # - visualize_variant_expression_effects: Show REF vs ALT variant effects on gene expression
+    # REMOVED - Too specific:
+    # - visualize_custom_annotations: Plot custom annotations like polyadenylation sites
+    - visualize_chromatin_accessibility: Visualize DNASE and ATAC chromatin accessibility
+    - visualize_splicing_effects: Visualize splicing predictions with SPLICE_SITES, SPLICE_SITE_USAGE, SPLICE_JUNCTIONS
+    # REMOVED - Redundant with visualize_variant_effects:
+    # - visualize_variant_splicing_effects: Show REF vs ALT variant effects on splicing with sashimi plots
+    - visualize_histone_modifications: Visualize CHIP_HISTONE predictions with custom colors
+    - visualize_tf_binding: Visualize CHIP_TF transcription factor binding predictions
+    - visualize_contact_maps: Visualize CONTACT_MAPS DNA-DNA contact predictions
+"""
+
+import sys
+from pathlib import Path
+from fastmcp import FastMCP
+import requests
+
+# Import the MCP tools from the tools folder
+from tools.batch_variant_scoring import batch_variant_scoring_mcp
+from tools.essential_commands import essential_commands_mcp
+
+# Define the MCP server
+mcp = FastMCP(name = "AlphaGenome")
+
+# Mount the tools
+mcp.mount(batch_variant_scoring_mcp)
+mcp.mount(essential_commands_mcp)
+
+# Run the MCP server
+if __name__ == "__main__":
+    mcp.run(show_banner=False)

README.md

@@ -1,10 +1,11 @@
 ---
-title: Alphagenome Agent
-emoji: 🏆
-colorFrom: green
-colorTo: red
+title: Alpha
+emoji: 👁
+colorFrom: blue
+python_version: 3.11
+colorTo: gray
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 5.49.0
 app_file: app.py
 pinned: false
 ---

agent_v3.py

@@ -0,0 +1,587 @@
+"""
+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, BaseMessage, HumanMessage, SystemMessage
+from langchain_openai import ChatOpenAI
+from dotenv import load_dotenv
+
+# Import core types and 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,
+    StateManager,
+    WorkflowEngine,
+    # Tools
+    ToolManager,
+    ToolSource,
+    ToolSelector,
+    # Execution
+    Timing,
+    PythonExecutor
+)
+
+# Load environment variables
+load_dotenv("./.env")
+
+
+def get_system_prompt(functions: Dict[str, dict], packages: Dict[str, str] = None) -> str:
+    """Generate system prompt using template and functions."""
+    if packages is None:
+        from core.constants import LIBRARY_CONTENT_DICT
+        packages = LIBRARY_CONTENT_DICT
+    return Template(SYSTEM_PROMPT_TEMPLATE).render(functions=functions, packages=packages)
+
+
+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,
+                 use_tool_manager: bool = True,
+                 use_tool_selection: bool = True):
+        """
+        Initialize the CodeAgent with unified tool management.
+
+        Args:
+            model: The language model to use for generation
+            config: Configuration for the agent
+            use_tool_manager: Whether to use the unified ToolManager (recommended)
+            use_tool_selection: Whether to use LLM-based tool selection (like Biomni)
+        """
+        self.model = model
+        self.config = config or AgentConfig()
+        self.use_tool_manager = use_tool_manager
+        self.use_tool_selection = use_tool_selection
+
+        # Cache selected tools to avoid re-selection at each step
+        self._selected_tools_cache = None
+
+        # Initialize modular components
+        self.package_manager = PackageManager()
+        self.console = ConsoleDisplay()
+        self.state_manager = StateManager()
+        self.plan_manager = PlanManager()
+
+        # Initialize unified tool management
+        if not self.use_tool_manager:
+            raise ValueError("ToolManager is required. Legacy mode (use_tool_manager=False) has been removed.")
+
+        self.tool_manager = ToolManager(self.console)
+
+        # Initialize tool selector for LLM-based tool selection
+        if self.use_tool_selection:
+            self.tool_selector = ToolSelector(self.model)
+        else:
+            self.tool_selector = None
+
+        # Initialize workflow engine
+        self.workflow_engine = WorkflowEngine(model, self.config, self.console, self.state_manager)
+
+        # 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."""
+
+        # Get all available tools first
+        all_schemas = self.tool_manager.get_tool_schemas(openai_format=True)
+        all_functions_dict = {schema['function']['name']: schema for schema in all_schemas}
+
+        # Use tool selection if enabled and not cached
+        if self.use_tool_selection and self.tool_selector and state.get("messages") and self._selected_tools_cache is None:
+            # Get the user's query from the first message
+            user_query = ""
+            for msg in state["messages"]:
+                if hasattr(msg, 'content') and msg.content:
+                    user_query = msg.content
+                    break
+
+            if user_query:
+                # Prepare tools for selection (convert schemas to tool info format)
+                available_tools = {}
+                for tool_name, schema in all_functions_dict.items():
+                    available_tools[tool_name] = {
+                        'description': schema['function'].get('description', 'No description'),
+                        'source': 'tool_manager'  # Could be enhanced to show actual source
+                    }
+
+                # Select relevant tools using LLM (only once)
+                selected_tool_names = self.tool_selector.select_tools_for_task(
+                    user_query, available_tools, max_tools=15
+                )
+
+                # Cache the selected tools
+                self._selected_tools_cache = {name: all_functions_dict[name]
+                                            for name in selected_tool_names
+                                            if name in all_functions_dict}
+
+                self.console.console.print(f"🎯 Selected {len(self._selected_tools_cache)} tools from {len(all_functions_dict)} available tools (cached for session)")
+                functions_dict = self._selected_tools_cache
+            else:
+                functions_dict = all_functions_dict
+        elif self.use_tool_selection and self._selected_tools_cache is not None:
+            # Use cached selected tools
+            functions_dict = self._selected_tools_cache
+        else:
+            # No tool selection or selection disabled
+            functions_dict = all_functions_dict
+
+        all_packages = self.package_manager.get_all_packages()
+        system_prompt = get_system_prompt(functions_dict, all_packages)
+
+        # Truncate conversation history to prevent context overflow
+        messages = [SystemMessage(content=system_prompt)] + state["messages"]
+
+        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 self.state_manager.create_state_dict(
+            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 self.plan_manager.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()
+
+                # Execute regular code in persistent environment (tools already injected)
+                result = self.python_executor(code)
+
+                # Include both the code and result in the observation
+                obs = f"\n<observation>\nCode Output:\n{result}</observation>"
+                return self.state_manager.create_state_dict(
+                    messages=[AIMessage(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 self.state_manager.create_state_dict(
+                    messages=[AIMessage(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 self.state_manager.create_state_dict(
+                messages=[AIMessage(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()
+
+    # ====================
+    # 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 list_mcp_tools(self) -> List[Dict]:
+        """List all loaded MCP tools."""
+        return self.tool_manager.list_tools(self.tool_manager.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)
+
+        # Initialize agent with functions using ToolManager
+        functions_dict = self.get_all_tool_functions()
+
+        # Display enhanced tool information
+        # Get detailed 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']['local'] > 0:
+            self.console.console.print(f"   📋 Local tools: {stats['by_source']['local']}")
+        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")
+
+        # Inject functions into Python executor
+        self.python_executor.send_functions(functions_dict)
+
+        # Import available packages using PackageManager
+        imported_packages, failed_packages = self.package_manager.import_packages(self.python_executor)
+        self.console.print_packages_info(imported_packages, failed_packages)
+
+        # Inject any initial variables
+        state_variables = {}
+        self.python_executor.send_variables(state_variables)
+
+        # Create initial state using StateManager
+        input_state = self.state_manager.create_state_dict(
+            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)
+
+        # Save trace and summary if requested
+        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__":
+    # Example usage of the fully modular CodeAgent architecture
+    # Create LLM client
+    model = ChatOpenAI(
+        model="google/gemini-2.5-flash",
+        base_url="https://openrouter.ai/api/v1",
+        temperature=0.7,
+        api_key=os.environ["OPENROUTER_API_KEY"],
+    )
+    
+    model = ChatAnthropic(model='claude-sonnet-4-5-20250929')
+    
+    
+
+    # Create configuration
+    config = AgentConfig(
+        max_steps=15,
+        max_conversation_length=30,
+        retry_attempts=3,
+        timeout_seconds=1200,
+        verbose=True
+    )
+
+    # Create agent with unified tool management and LLM-based tool selection
+    agent = CodeAgent(model=model, config=config, use_tool_manager=True, use_tool_selection=True)
+
+    # Demonstrate tool management capabilities
+    print("\n🔧 Tool Management Demo:")
+
+    # Show tool statistics
+    stats = agent.get_tool_statistics()
+    print(f"📊 Tool Statistics: {stats}")
+
+    # Add MCP tools
+    try:
+        print("🔧 Loading MCP tools...")
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        config_path = os.path.join(script_dir, "mcp_config.yaml")
+        agent.add_mcp(config_path)
+        print("✅ MCP tools loaded successfully!")
+
+        # Show enhanced MCP status
+        agent.show_mcp_status()
+
+        # Show all available tools by source
+        all_tools = agent.list_tools("all", include_details=True)
+        print(f"\n📋 All tools loaded:")
+        for tool in all_tools:
+            print(f"   • {tool['name']} ({tool['source']}) - {tool['description'][:50]}...")
+
+    except Exception as e:
+        print(f"⚠️ Could not load MCP tools: {e}")
+
+    # Validate tools
+    issues = agent.validate_tools()
+    if any(issues.values()):
+        print(f"⚠️  Tool validation issues: {issues}")
+    else:
+        print("✅ All tools validated successfully!")
+    # List available packages
+    print(f"\n📦 Available packages: {list(agent.get_all_packages().keys())}")
+    # Run the agent with a query and save trace/summary
+    print("\n🚀 Running agent with trace and summary saving...")
+    result = agent.run(
+        query="""
+        Use AlphaGenome MCP to analyze heart gene expression data to identify the causal gene
+        for the variant chr11:116837649:T>G, associated with Hypoalphalipoproteinemia. My API key is: AIzaSyD1USDNy9WqfIROICB3FWI1wJHmkO2z21U.
+        """,
+        save_trace=True,      # Save complete execution trace
+        save_summary=True,    # Save execution summary
+        trace_dir="traces"    # Directory to save files
+    )
+
+    # You can also access trace and summary programmatically
+    print("\n📊 Execution Summary:")
+    summary = agent.get_summary()
+    print(f"   Total steps: {summary.get('total_steps', 0)}")
+    print(f"   Code executions: {len(summary.get('code_executions', []))}")
+    print(f"   Observations: {len(summary.get('observations', []))}")
+    print(f"   Errors: {len(summary.get('errors', []))}")
+
+    # You can save trace/summary manually after execution
+    # agent.save_trace("custom_trace.json")
+    # agent.save_summary("custom_summary.json")

app.py

@@ -1,69 +1,656 @@
+"""
+Gradio Interface for LangGraph ReAct Agent
+A production-ready web interface with real-time streaming output
+Styled similar to gradio_ui.py with sidebar layout
+"""
+
+import os
+import re
+import sys
+import time
+from typing import Generator, Optional
+from dataclasses import dataclass
+
 import gradio as gr
+from gradio.themes.utils import fonts
+from dotenv import load_dotenv
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_anthropic import ChatAnthropic
 
-embed_html1 = """
-<iframe width="1120" height="620" 
-src="https://www.youtube.com/embed/WmSC6zi8CIY?si=8h9HR1r9bZIpOaVN" 
-title="YouTube video player" frameborder="0" 
-allow="accelerometer; autoplay; clipboard-write; encrypted-media; 
-gyroscope; picture-in-picture; web-share" 
-referrerpolicy="strict-origin-when-cross-origin" 
-allowfullscreen></iframe>
-"""
+# Import the agent from agent_v3
+sys.path.append(os.path.dirname(__file__))
+from agent_v3 import CodeAgent
+from core.types import AgentConfig
 
-embed_html2 = """
-<iframe width="1120" height="620" 
-src="https://www.youtube.com/embed/QMNqvybf_9U?si=Y1AAy5V_S7Nht4Vs" 
-title="YouTube video player" frameborder="0" 
-allow="accelerometer; autoplay; clipboard-write; encrypted-media; 
-gyroscope; picture-in-picture; web-share" 
-referrerpolicy="strict-origin-when-cross-origin" 
-allowfullscreen></iframe>
-"""
+# Import support modules
+from managers import PythonExecutor
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-# Demos
-Below, we showcase demos of AI agents created by Paper2Agent, illustrating how each agent applies the tools from its source paper to tackle scientific tasks.
+# Load environment variables
+load_dotenv("./.env")
 
-## 🧬 AlphaGenome Agent for Genomic Data Interpretation
-**Example query:**
-```
-Analyze heart gene expression data with AlphaGenome MCP to identify the causal genefor the variant chr11:116837649:T>G, associated with Hypoalphalipoproteinemia.
-```
-        """
-    )
-    gr.HTML(embed_html1)
-    gr.Markdown(
-        """
-# 🤖 How to Create AlphaGenome and other Paper Agents?
-
-To streamline usage, we recommend creating Paper Agents by connecting Paper MCP servers to an AI coding agent, such as [Claude Code](https://www.anthropic.com/claude-code) or the [Google Gemini CLI](https://google-gemini.github.io/gemini-cli/) (it's free with a Google account!).
-
-We are also actively developing our own Chatbot-based agent, which will be released soon.
-
-## ⚙️ Using Claude Code
-First, install and set up Claude Code:
-```bash
-npm install -g @anthropic-ai/claude-code
-claude
-```
-After setup, link Claude Code with the Paper MCP server of interest.
-For example, to create an AlphaGenome Agent, run:
-```
-claude mcp add \
---transport http \
-alphagenome \
-https://Paper2Agent-alphagenome-mcp.hf.space/mcp
 
-claude
+class GradioAgentUI:
+    """
+    Gradio interface for interacting with the LangGraph ReAct Agent.
+    Styled similar to the smolagents GradioUI with sidebar layout.
+    """
+    
+    def __init__(self, model=None, config=None):
+        """Initialize the Gradio Agent UI."""
+
+        # Default model
+        if model is None:
+            api_key = os.environ.get("OPENROUTER_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "OPENROUTER_API_KEY environment variable is not set. "
+                    "Please set it or provide a model instance."
+                )
+            # model = ChatOpenAI(
+            #     model="google/gemini-2.5-flash",
+            #     base_url="https://openrouter.ai/api/v1",
+            #     temperature=0.7,
+            #     api_key=api_key,
+            # )
+            
+            model = ChatAnthropic(
+                model='claude-sonnet-4-5-20250929',
+                temperature=0.7,
+                api_key='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'
+            )
+            
+            
+            
+
+        # Default config
+        if config is None:
+            config = AgentConfig(
+                max_steps=15,
+                max_conversation_length=30,
+                retry_attempts=3,
+                timeout_seconds=1200,
+                verbose=True
+            )
+
+        self.model = model
+        self.config = config
+        self.name = "AlphaGenome Agent Interface"
+        self.description = "AlphaGenome Agent is automatically created by [Paper2Agent](https://github.com/jmiao24/Paper2Agent) to use [AlphaGenome](https://github.com/google-deepmind/alphagenome) and interpret DNA variants."
+
+    def get_step_footnote(self, step_num: int, duration: float) -> str:
+        """Create a footnote for a step with timing information."""
+        return f'<span style="color: #888; font-size: 0.9em;">Step {step_num} | Duration: {duration:.2f}s</span>'
+    
+    def format_code_block(self, code: str) -> str:
+        """Format code as a Python code block."""
+        code = code.strip()
+        if not code.startswith("```"):
+            code = f"```python\n{code}\n```"
+        return code
+    
+    def extract_content_parts(self, message_content: str) -> dict:
+        """Extract different parts from the message content."""
+        parts = {
+            'thinking': '',
+            'plan': None,
+            'code': None,
+            'solution': None,
+            'error': None,
+            'observation': None
+        }
+        
+        # Extract thinking (content before any tags)
+        thinking_pattern = r'^(.*?)(?=<(?:execute|solution|error|observation)|$)'
+        thinking_match = re.match(thinking_pattern, message_content, re.DOTALL)
+        if thinking_match:
+            thinking = thinking_match.group(1).strip()
+            # Remove plan from thinking
+            plan_pattern = r'\d+\.\s*\[[^\]]*\]\s*[^\n]+(?:\n\d+\.\s*\[[^\]]*\]\s*[^\n]+)*'
+            thinking = re.sub(plan_pattern, '', thinking).strip()
+            # Remove any existing "Thinking:" or "Plan:" labels (including duplicates)
+            thinking = re.sub(r'(^|\n)(Thinking:|Plan:)\s*', '\n', thinking).strip()
+            # Remove any duplicate Thinking: that might remain
+            thinking = re.sub(r'Thinking:\s*', '', thinking).strip()
+            if thinking:
+                parts['thinking'] = thinking
+        
+        # Extract plan
+        plan_pattern = r'\d+\.\s*\[[^\]]*\]\s*[^\n]+(?:\n\d+\.\s*\[[^\]]*\]\s*[^\n]+)*'
+        plan_match = re.search(plan_pattern, message_content)
+        if plan_match:
+            parts['plan'] = plan_match.group(0)
+        
+        # Extract code
+        code_match = re.search(r'<execute>(.*?)</execute>', message_content, re.DOTALL)
+        if code_match:
+            parts['code'] = code_match.group(1).strip()
+        
+        # Extract solution
+        solution_match = re.search(r'<solution>(.*?)</solution>', message_content, re.DOTALL)
+        if solution_match:
+            parts['solution'] = solution_match.group(1).strip()
+        
+        # Extract error
+        error_match = re.search(r'<error>(.*?)</error>', message_content, re.DOTALL)
+        if error_match:
+            parts['error'] = error_match.group(1).strip()
+        
+        # Extract observation
+        obs_match = re.search(r'<observation>(.*?)</observation>', message_content, re.DOTALL)
+        if obs_match:
+            parts['observation'] = obs_match.group(1).strip()
+        
+        return parts
+    
+    def truncate_output(self, text: str, max_lines: int = 20) -> str:
+        """Truncate output to specified number of lines."""
+        lines = text.split('\n')
+        if len(lines) > max_lines:
+            truncated = '\n'.join(lines[:max_lines])
+            truncated += f"\n... (truncated {len(lines) - max_lines} lines)"
+            return truncated
+        return text
+    
+    def stream_agent_response(self, agent: CodeAgent, query: str) -> Generator:
+        """Stream agent responses as Gradio ChatMessages with structured blocks."""
+
+        # Get all tool functions using the agent's method
+        tools_dict = agent.get_all_tool_functions()
+        agent.python_executor.send_functions(tools_dict)
+        agent.python_executor.send_variables({})
+
+        # Create initial state
+        input_state = {
+            "messages": [HumanMessage(content=query)],
+            "step_count": 0,
+            "error_count": 0,
+            "start_time": time.time(),
+            "current_plan": None
+        }
+
+        # Track state
+        displayed_reasoning = set()
+        previous_plan = None
+        step_timings = {}
+
+        try:
+            # Stream through the workflow execution
+            for state in agent.workflow_engine.graph.stream(input_state, stream_mode="values"):
+                step_count = state.get("step_count", 0)
+                error_count = state.get("error_count", 0)
+                current_plan = state.get("current_plan")
+
+                # Track timing
+                if step_count not in step_timings:
+                    step_timings[step_count] = time.time()
+
+                message = state["messages"][-1]
+
+                if isinstance(message, AIMessage):
+                    content = message.content
+                    parts = self.extract_content_parts(content)
+
+                    # First yield step header only
+                    if step_count > 0:
+                        step_header = f"## Step {step_count}\n"
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=step_header,
+                            metadata={"status": "done"}
+                        )
+                    
+                    # Display Thinking block (styled like Current Plan)
+                    if parts['thinking'] and len(parts['thinking']) > 20:
+                        thinking_text = parts['thinking']
+                        # Clean up any remaining labels
+                        thinking_text = re.sub(r'(Thinking:|Plan:)\s*', '', thinking_text).strip()
+                        # Clean up excessive whitespace - replace multiple newlines with max 2
+                        thinking_text = re.sub(r'\n\s*\n\s*\n+', '\n\n', thinking_text).strip()
+                        # Also clean up any leading/trailing whitespace on each line
+                        thinking_text = '\n'.join(line.strip() for line in thinking_text.split('\n') if line.strip())
+
+                        content_hash = hash(thinking_text)
+                        if content_hash not in displayed_reasoning and thinking_text:
+                            thinking_block = f"""<div style="background-color: #f8f9fa; border-left: 4px solid #6b7280; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #374151; margin-bottom: 8px;">🤔 Thinking</div>
+<div style="margin: 0; white-space: pre-line; font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; color: #000000; background-color: #ffffff; padding: 10px; border-radius: 4px; font-size: 14px; line-height: 1.6;">{thinking_text}</div>
+</div>"""
+                            yield gr.ChatMessage(
+                                role="assistant",
+                                content=thinking_block,
+                                metadata={"status": "done"}
+                            )
+                            displayed_reasoning.add(content_hash)
+
+                    # Then display Current Plan after thinking
+                    if current_plan and current_plan != previous_plan:
+                        formatted_plan = current_plan.replace('[ ]', '☐').replace('[✓]', '✅').replace('[✗]', '❌')
+
+                        plan_block = f"""<div style="background-color: #f8f9fa; border-left: 4px solid #000000; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #000000; margin-bottom: 8px;">📋 Current Plan</div>
+<pre style="margin: 0; white-space: pre-wrap; font-family: 'Consolas', 'Monaco', 'Courier New', monospace; color: #000000; background-color: #ffffff; padding: 10px; border-radius: 4px; font-size: 14px; line-height: 1.6;">{formatted_plan}</pre>
+</div>"""
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=plan_block,
+                            metadata={"status": "done"}
+                        )
+                        previous_plan = current_plan
+
+                    # Display Executing Code as independent block with proper syntax highlighting
+                    if parts['code']:
+                        # Format as markdown code block for proper rendering
+                        code_block = f"""<div style="background-color: #f7fafc; border-left: 4px solid #48bb78; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #22543d; margin-bottom: 8px;">⚡ Executing Code</div>
+</div>
+
+```python
+{parts['code']}
+```"""
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=code_block,
+                            metadata={"status": "done"}
+                        )
+
+                    # Display Execution Result as independent block
+                    if parts['observation']:
+                        truncated = self.truncate_output(parts['observation'])
+
+                        result_block = f"""<div style="background-color: #fef5e7; border-left: 4px solid #f6ad55; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #744210; margin-bottom: 8px;">📊 Execution Result</div>
+</div>
+
 ```
-✅ You will now have an **AlphaGenome Agent** ready for genomics data interpretation.
-
-# 🔗 Connectable Paper MCP Servers
-* AlphaGenome: https://Paper2Agent-alphagenome-mcp.hf.space/mcp
-* Scanpy: https://Paper2Agent-scanpy-mcp.hf.space/mcp
-* TISSUE: https://Paper2Agent-tissue-mcp.hf.space/mcp
-        """
-     )
-demo.launch()
+{truncated}
+```"""
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=result_block,
+                            metadata={"status": "done"}
+                        )
+
+                    # Display solution with special formatting as independent block
+                    if parts['solution']:
+                        solution_block = f"""<div style="background-color: #d1fae5; border-left: 4px solid #10b981; padding: 16px; margin: 10px 0; border-radius: 6px;">
+<div style="font-weight: bold; color: #065f46; margin-bottom: 12px; font-size: 16px;">✅ Final Solution</div>
+<div class="solution-content" style="color: #1f2937 !important; background-color: #ffffff; padding: 16px; border-radius: 4px; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif; font-size: 15px; line-height: 1.7; border: 1px solid #e5e7eb;">
+
+{parts['solution']}
+
+</div>
+<style>
+.solution-content, .solution-content * {{
+    color: #1f2937 !important;
+    background-color: transparent !important;
+}}
+.solution-content code, .solution-content pre {{
+    background-color: #f3f4f6 !important;
+    color: #1f2937 !important;
+    padding: 2px 6px;
+    border-radius: 3px;
+    font-family: 'Monaco', 'Menlo', 'Courier New', monospace;
+}}
+.solution-content pre {{
+    padding: 12px;
+    overflow-x: auto;
+}}
+.solution-content table {{
+    border-collapse: collapse;
+    width: 100%;
+    margin: 10px 0;
+    background-color: #ffffff !important;
+}}
+.solution-content td, .solution-content th {{
+    border: 1px solid #d1d5db;
+    padding: 8px;
+    color: #1f2937 !important;
+    background-color: #ffffff !important;
+}}
+.solution-content th {{
+    background-color: #f3f4f6 !important;
+    font-weight: 600;
+}}
+.solution-content p, .solution-content div, .solution-content span {{
+    background-color: transparent !important;
+}}
+</style>
+</div>"""
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=solution_block,
+                            metadata={"status": "done"}
+                        )
+
+                    # Display error with special formatting
+                    if parts['error']:
+                        error_block = f"""<div style="background-color: #fed7d7; border-left: 4px solid #fc8181; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #742a2a; margin-bottom: 8px;">⚠️ Error</div>
+<div style="color: #742a2a;">{parts['error']}</div>
+</div>"""
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=error_block,
+                            metadata={"status": "done"}
+                        )
+
+                    # Add step footnote with timing ONLY after Execution Result
+                    if parts['observation'] and step_count in step_timings:
+                        duration = time.time() - step_timings[step_count]
+                        footnote = f'<div style="color: #718096; font-size: 0.875em; margin-top: 8px;">Step {step_count} | Duration: {duration:.2f}s</div>'
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content=footnote,
+                            metadata={"status": "done"}
+                        )
+
+                    # Add separator between steps
+                    if step_count > 0:
+                        yield gr.ChatMessage(
+                            role="assistant",
+                            content='<hr style="border: none; border-top: 1px solid #e2e8f0; margin: 20px 0;">',
+                            metadata={"status": "done"}
+                        )
+        
+        except Exception as e:
+            error_block = f"""<div style="background-color: #fed7d7; border-left: 4px solid #fc8181; padding: 12px; margin: 10px 0; border-radius: 4px;">
+<div style="font-weight: bold; color: #742a2a; margin-bottom: 8px;">💥 Critical Error</div>
+<div style="color: #742a2a;">Error during agent execution: {str(e)}</div>
+</div>"""
+            yield gr.ChatMessage(
+                role="assistant",
+                content=error_block,
+                metadata={"status": "done"}
+            )
+    
+    def interact_with_agent(self, prompt: str, api_key: str, messages: list, session_state: dict) -> Generator:
+        """Handle interaction with the agent."""
+
+        # Store original user input for display
+        original_prompt = prompt
+
+        # Add API key to the prompt if provided (for agent processing)
+        if api_key and api_key.strip():
+            prompt = f"{prompt} My API key is: {api_key.strip()}."
+
+        # Initialize agent in session if needed
+        if "agent" not in session_state:
+            session_state["agent"] = CodeAgent(
+                model=self.model,
+                config=self.config,
+                use_tool_manager=True,
+                use_tool_selection=False  # Disable for Gradio UI by default
+            )
+
+            # Try to load MCP config if available
+            mcp_config_path = "./mcp_config.yaml"
+            if os.path.exists(mcp_config_path):
+                try:
+                    session_state["agent"].add_mcp(mcp_config_path)
+                    print(f"✅ Loaded MCP tools from {mcp_config_path}")
+                except Exception as e:
+                    print(f"⚠️ Could not load MCP tools: {e}")
+
+        agent = session_state["agent"]
+
+        try:
+            # Add user message immediately (use original prompt for display)
+            messages.append(gr.ChatMessage(role="user", content=original_prompt, metadata={"status": "done"}))
+            yield messages
+
+            # Add a "thinking" indicator immediately
+            messages.append(gr.ChatMessage(role="assistant", content="🤔 Processing...", metadata={"status": "pending"}))
+            yield messages
+
+            # Remove the thinking indicator before adding real content
+            messages.pop()
+
+            # Stream agent responses with real-time updates
+            for msg in self.stream_agent_response(agent, prompt):
+                messages.append(msg)
+                yield messages
+
+        except Exception as e:
+            messages.append(
+                gr.ChatMessage(
+                    role="assistant",
+                    content=f"Error: {str(e)}",
+                    metadata={"title": "💥 Error", "status": "done"}
+                )
+            )
+            yield messages
+    
+    def create_app(self):
+        """Create the Gradio app with sidebar layout."""
+        
+        # Create custom theme with modern fonts
+        modern_theme = gr.themes.Monochrome(
+            font=fonts.GoogleFont("Inter"),
+            font_mono=fonts.GoogleFont("JetBrains Mono")
+        )
+        
+        with gr.Blocks(theme=modern_theme, fill_height=True, title=self.name, css="""
+            /* Hide Gradio footer */
+            .footer {display: none !important;}
+            footer {display: none !important;}
+            .gradio-footer {display: none !important;}
+            #footer {display: none !important;}
+            [class*="footer"] {display: none !important;}
+            [id*="footer"] {display: none !important;}
+            .block.svelte-1scc9gv {display: none !important;}
+            .built-with-gradio {display: none !important;}
+            .gradio-container footer {display: none !important;}
+            """) as demo:
+            # Force light theme and hide footer
+            demo.load(js="""
+            () => {
+                document.body.classList.remove('dark');
+                document.querySelector('gradio-app').classList.remove('dark');
+                const url = new URL(window.location);
+                url.searchParams.set('__theme', 'light');
+                window.history.replaceState({}, '', url);
+
+                // Hide footer elements
+                setTimeout(() => {
+                    const footers = document.querySelectorAll('footer, .footer, .gradio-footer, #footer, [class*="footer"], [id*="footer"], .built-with-gradio');
+                    footers.forEach(footer => footer.style.display = 'none');
+
+                    // Also hide any Gradio logo/branding
+                    const brandingElements = document.querySelectorAll('a[href*="gradio"], .gradio-logo, [alt*="gradio"]');
+                    brandingElements.forEach(el => el.style.display = 'none');
+                }, 100);
+            }
+            """)
+            # Session state
+            session_state = gr.State({})
+            stored_messages = gr.State([])
+            
+            with gr.Row():
+                # Sidebar
+                with gr.Column(scale=1):
+                    gr.Markdown(f"# {self.name}")
+                    gr.Markdown(f"> {self.description}")
+                    
+                    gr.Markdown("---")
+                    
+                    # API Key section
+                    with gr.Group():
+                        gr.Markdown("**AlphaGenome API Key**")
+                        api_key_input = gr.Textbox(
+                            label="API Key",
+                            type="password",
+                            placeholder="Enter your AlphaGenome API key",
+                            value="AIzaSyD1USDNy9WqfIROICB3FWI1wJHmkO2z21U"
+                        )
+
+                    # Input section
+                    with gr.Group():
+                        gr.Markdown("**Your Request**")
+                        text_input = gr.Textbox(
+                            lines=4,
+                            label="Query",
+                            placeholder="Enter your query here and press Enter or click Submit",
+                            value="""Use AlphaGenome MCP to analyze heart gene expression data to identify the causal gene
+for the variant chr11:116837649:T>G, associated with Hypoalphalipoproteinemia."""
+                        )
+                        submit_btn = gr.Button("🚀 Submit", variant="primary", size="lg")
+                    
+                    # Configuration section
+                    with gr.Accordion("⚙️ Configuration", open=False):
+                        max_steps_input = gr.Slider(
+                            label="Max Steps",
+                            minimum=5,
+                            maximum=50,
+                            value=self.config.max_steps,
+                            step=1
+                        )
+
+                        temperature_input = gr.Slider(
+                            label="Temperature",
+                            minimum=0.0,
+                            maximum=1.0,
+                            value=0.7,
+                            step=0.1
+                        )
+
+                        apply_config_btn = gr.Button("Apply Configuration", size="sm")
+                    
+                    # Example queries
+                    with gr.Accordion("📚 Example Queries", open=False):
+                        gr.Examples(
+                            examples=[
+                                ["What's the quantile score of chr3:197081044:TACTC>T on splice junction in Artery (tibial)?"],
+                                ["What are the raw and quantile scores of the variant chr3:120280774:G>T for chromatin accessibility in the GM12878 cell line?"],
+                            ],
+                            inputs=text_input,
+                        )
+                    
+                    gr.Markdown("---")
+                    # gr.HTML(
+                    #     "<center><small>Powered by LangGraph & OpenRouter</small></center>"
+                    # )
+                
+                # Main chat area
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(
+                        label="Agent Conversation",
+                        type="messages",
+                        height=700,
+                        show_copy_button=True,
+                        avatar_images=(
+                            None,  # Default user avatar
+                            "images.png"   # Assistant avatar
+                        ),
+                        latex_delimiters=[
+                            {"left": r"$$", "right": r"$$", "display": True},
+                            {"left": r"$", "right": r"$", "display": False},
+                        ],
+                        render_markdown=True,
+                        sanitize_html=False,  # Allow custom HTML for better formatting
+                    )
+                    
+                    with gr.Row():
+                        clear_btn = gr.Button("🗑️ Clear", size="sm")
+                        # Note: stop_btn is created but not wired up yet
+                        # This could be used for future cancellation functionality
+                        stop_btn = gr.Button("⏹️ Stop", size="sm", variant="stop")
+            
+            # Event handlers
+            def update_config(max_steps, temperature, session_state):
+                """Update agent configuration."""
+                if "agent" in session_state:
+                    agent = session_state["agent"]
+                    agent.config.max_steps = max_steps
+                    if hasattr(agent.model, 'temperature'):
+                        agent.model.temperature = temperature
+                return "Configuration updated!"
+            
+            def clear_chat():
+                """Clear the chat."""
+                return [], []
+            
+            # Wire up events
+            text_input.submit(
+                lambda x: ("", gr.Button(interactive=False)),
+                [text_input],
+                [text_input, submit_btn]
+            ).then(
+                self.interact_with_agent,
+                [stored_messages, api_key_input, chatbot, session_state],
+                [chatbot]
+            ).then(
+                lambda: gr.Button(interactive=True),
+                None,
+                [submit_btn]
+            )
+            
+            submit_btn.click(
+                lambda x: (x, "", gr.Button(interactive=False)),
+                [text_input],
+                [stored_messages, text_input, submit_btn]
+            ).then(
+                self.interact_with_agent,
+                [stored_messages, api_key_input, chatbot, session_state],
+                [chatbot]
+            ).then(
+                lambda: gr.Button(interactive=True),
+                None,
+                [submit_btn]
+            )
+            
+            apply_config_btn.click(
+                update_config,
+                [max_steps_input, temperature_input, session_state],
+                None
+            )
+            
+            clear_btn.click(
+                clear_chat,
+                None,
+                [chatbot, stored_messages]
+            )
+        
+        return demo
+    
+    def launch(self, share: bool = False, **kwargs):
+        """Launch the Gradio app."""
+        app = self.create_app()
+        # Set defaults if not provided in kwargs
+        kwargs.setdefault('server_name', '0.0.0.0')
+        kwargs.setdefault('server_port', 7860)
+        kwargs.setdefault('show_error', True)
+
+        app.queue(max_size=10).launch(
+            share=share,
+            show_api=False,          # 隐藏 "Use via API" 链接
+            favicon_path=None,       # 移除默认favicon
+            **kwargs
+        )
+
+
+if __name__ == "__main__":
+    # Check for API key before starting
+    if not os.environ.get("OPENROUTER_API_KEY"):
+        print("\n⚠️  Error: OPENROUTER_API_KEY environment variable is not set!")
+        print("\nPlease set it using one of these methods:")
+        print("1. Export it in your shell:")
+        print("   export OPENROUTER_API_KEY='your-api-key-here'")
+        print("\n2. Create a .env file in the current directory with:")
+        print("   OPENROUTER_API_KEY=your-api-key-here")
+        print("\n3. Or provide your own model instance when initializing GradioAgentUI")
+        sys.exit(1)
+
+    try:
+        # Create and launch the interface with optional MCP config
+        ui = GradioAgentUI()
+
+        # Optional: Load MCP tools if config exists
+        mcp_config_path = "./mcp_config.yaml"
+        if os.path.exists(mcp_config_path):
+            print(f"Found MCP config at {mcp_config_path}")
+
+        ui.launch(share=False, quiet=False)
+    except Exception as e:
+        print(f"\n❌ Error starting Gradio interface: {e}")
+        sys.exit(1)

core/__init__.py

@@ -0,0 +1,8 @@
+"""
+Core module for CodeAct agent - contains shared types and utilities.
+"""
+
+from .types import AgentState, AgentConfig
+from .constants import LIBRARY_CONTENT_DICT, SYSTEM_PROMPT_TEMPLATE
+
+__all__ = ['AgentState', 'AgentConfig', 'LIBRARY_CONTENT_DICT', 'SYSTEM_PROMPT_TEMPLATE']

core/constants.py

@@ -0,0 +1,259 @@
+"""
+Constants and templates for CodeAct agent.
+"""
+
+# Available packages with descriptions
+LIBRARY_CONTENT_DICT = {
+    "numpy": "[Python Package] The fundamental package for scientific computing with Python, providing support for arrays, matrices, and mathematical functions.",
+    "scipy": "[Python Package] A Python library for scientific and technical computing, including modules for optimization, linear algebra, integration, and statistics.",
+}
+
+# System prompt template
+SYSTEM_PROMPT_TEMPLATE = """You are an expert assistant who can solve any task using Python code execution. You will be given a task to solve as best you can.
+
+## Planning & Execution Workflow
+
+### Step 1: Make a Plan
+For every task, start by creating a clear, numbered plan with checkboxes:
+```
+1. [ ] First step
+2. [ ] Second step
+3. [ ] Third step
+```
+
+### Step 2: Execute & Update
+After completing each step, update the checklist with [✓]:
+```
+1. [✓] First step (completed)
+2. [ ] Second step
+3. [ ] Third step
+```
+
+If a step fails, mark it with [✗] and add a corrective step:
+```
+1. [✓] First step (completed)
+2. [✗] Second step (failed because...)
+3. [ ] Corrected second step
+4. [ ] Third step
+```
+
+**Always show the updated plan after each step** so progress can be tracked.
+
+## Response Format
+
+At each turn, first provide your **thinking and reasoning** based on the conversation history.
+
+Then choose ONE of these two options:
+
+### Option 1: Execute Code (<execute> tag)
+Use this to run Python code and get results. The output will appear in <observation></observation> tags.
+
+**Format:**
+```
+<execute>
+print("Hello World!")
+</execute>
+```
+
+**CRITICAL**: Always end with `</execute>` tag.
+
+### Option 2: Provide Solution (<solution> tag)
+Use this when you have the final answer ready to present to the user.
+
+**Format:**
+```
+The answer is <solution> Your final answer </solution>
+```
+
+**CRITICAL**:
+- Always end with `</solution>` tag.
+- **Always provide a comprehensive summary** in your final solution that includes:
+  - A clear statement of what was accomplished
+  - Key results, findings, or outputs
+  - Any important data, numbers, or visualizations generated
+  - Steps taken to reach the solution (brief overview)
+  - Any files created or modified
+  - Next steps or recommendations (if applicable)
+
+## Code Execution Best Practices
+
+### Printing Results
+**YOU MUST print outputs** - the system cannot see unpublished results:
+```python
+# ✓ CORRECT
+result = some_function(param1, param2)
+print(result)
+
+# ✗ WRONG - output is invisible
+some_function(param1, param2)
+```
+
+### Working with Data
+When working with data structures, always inspect them **but limit output to avoid context overflow**:
+```python
+# For DataFrames - ALWAYS use head() for large tables
+print(f"Shape: {df.shape}")
+print(f"Columns: {list(df.columns)}")
+print(df.head(10))  # Show first 10 rows max
+
+# For large DataFrames, also show basic info
+if len(df) > 20:
+    print(f"DataFrame has {len(df)} rows. Showing first 10 rows only.")
+
+# For lists/dicts
+print(f"Length: {len(data)}")
+print(data[:5])  # Show first few items
+
+# For any large output, use head() or limit display
+if hasattr(data, 'head'):
+    print(data.head(10))  # For pandas objects
+else:
+    print(str(data)[:1000] + "..." if len(str(data)) > 1000 else data)
+```
+
+### Function Calls
+When calling available functions, **capture and print the output**:
+```python
+result = function_name(param1="value1", param2="value2")
+print(f"Function result: {result}")
+```
+
+### Creating Input Files for Functions
+If a function requires a file as input (e.g., CSV, Excel, TSV), **you must create the file first** with the exact columns specified in the function's parameter description:
+
+```python
+# Example: If function requires a CSV with specific columns
+import pandas as pd
+
+# Create DataFrame with required columns (as specified in function parameters)
+df = pd.DataFrame({
+    'column1': [...],  # Fill with appropriate data
+    'column2': [...],  # Fill with appropriate data
+    'column3': [...]   # Fill with appropriate data
+})
+
+# Save to file
+df.to_csv('input_file.csv', index=False)
+print(f"Created input file with shape: {df.shape}")
+
+# Then call the function
+result = function_name(file_path='input_file.csv')
+print(f"Function result: {result}")
+```
+
+**Important**: Always check the function's parameter descriptions to understand:
+- What columns are required in the input file
+- The expected format of each column
+- Any specific data types or constraints
+
+### Code Quality
+- Keep code simple and readable
+- Add comments for clarity
+- Print intermediate results to create a clear execution log
+- Handle errors gracefully
+
+## Important Rules
+
+1. **Every response must include EXACTLY ONE tag**: Either `<execute>` or `<solution>`
+2. **Never use both tags in the same response**
+3. **Never respond without any tags**
+4. **Always close tags properly** with `</execute>` or `</solution>`
+5. **Always print outputs** when executing code - unpublished results are invisible to the system
+
+## Examples
+
+### Example 1: Executing Code
+
+**User Query**: "Calculate the factorial of 10"
+
+**Correct Response**:
+```
+I'll calculate the factorial of 10 using Python.
+
+<execute>
+import math
+result = math.factorial(10)
+print(f"The factorial of 10 is: {result}")
+</execute>
+```
+
+**WRONG Response** (missing tags):
+```
+I'll calculate the factorial of 10.
+import math
+result = math.factorial(10)
+print(result)
+```
+
+### Example 2: Using Available Functions
+
+**User Query**: "Add 5 and 3"
+
+**Correct Response**:
+```
+I'll use the add_numbers function to add 5 and 3.
+
+<execute>
+result = add_numbers(5, 3)
+print(f"Result: {result}")
+</execute>
+```
+
+### Example 3: Providing Final Solution
+
+**User Query**: "What is 2+2?"
+
+**Correct Response**:
+```
+<execute>
+result = 2 + 2
+print(f"2 + 2 = {result}")
+</execute>
+```
+
+Then after seeing the output:
+```
+<solution>
+The answer is 4. I calculated 2 + 2 and got the result of 4.
+</solution>
+```
+
+**CRITICAL REMINDER**:
+- ALWAYS wrap code in `<execute>` tags
+- ALWAYS wrap final answers in `<solution>` tags
+- NEVER output code without `<execute>` tags
+
+## Available Functions
+
+You have access to the following functions. These functions are already available in your Python environment and can be called directly:
+
+{% for func_name, schema in functions.items() %}
+**{{ schema.function.name }}({% for param_name in schema.function.parameters.properties.keys() %}{{ param_name }}{{ ", " if not loop.last }}{% endfor %})**
+- Description: {{ schema.function.description }}
+- Parameters:
+  {% for param_name, param_info in schema.function.parameters.properties.items() %}
+  - {{ param_name }} ({{ param_info.type }}{% if param_info.enum %}, choices: {{ param_info.enum }}{% endif %}): {{ param_info.description }}{% if 'default' in param_info %} [default: {{ param_info.default }}]{% endif %}{% if param_name in schema.function.parameters.required %} [REQUIRED]{% endif %}
+  {% endfor %}
+- Example: result = {{ schema.function.name }}({% for param_name in schema.function.parameters.properties.keys() %}{{ param_name }}=example_value{{ ", " if not loop.last }}{% endfor %})
+{% if schema.function.parameters.required %}
+- **Important**: This function requires the following parameters: {{ schema.function.parameters.required|join(', ') }}
+{% endif %}
+
+{% endfor %}
+
+Important: These functions are pre-loaded in your Python environment. Call them directly like regular Python functions.
+**CRITICAL**: Always provide ALL required parameters when calling functions. Check the [REQUIRED] markers above.
+
+## Available Python Packages
+
+The following Python packages are available in your environment and can be imported directly:
+
+{% for package_name, description in packages.items() %}
+**{{ package_name }}**
+- {{ description }}
+- Import: `import {{ package_name }}`
+
+{% endfor %}
+
+Remember to import these packages before using them in your code.
+"""

core/types.py

@@ -0,0 +1,29 @@
+"""
+Shared type definitions for CodeAct agent.
+"""
+
+from typing import Annotated, List, Optional, TypedDict
+from dataclasses import dataclass
+from langchain_core.messages import BaseMessage
+from langgraph.graph.message import add_messages
+
+
+@dataclass
+class AgentConfig:
+    """Configuration for the Code Agent."""
+    max_steps: int = 20
+    max_conversation_length: int = 50
+    retry_attempts: int = 3
+    timeout_seconds: int = 900
+    verbose: bool = True
+    memory_window: int = 15  # Keep last N messages for context
+
+
+class AgentState(TypedDict):
+    """State type for the agent workflow."""
+    # reducer tells LangGraph to append new messages
+    messages: Annotated[List[BaseMessage], add_messages]
+    step_count: int
+    error_count: int
+    start_time: float
+    current_plan: Optional[str]

managers/__init__.py

@@ -0,0 +1,41 @@
+"""
+Manager modules for CodeAct agent - organized by subsystem.
+
+Subsystems:
+- execution: Python execution and monitoring
+- tools: Tool management, registry, and integrations
+- workflow: Workflow execution, state, and plan management
+- support: Console display and package management
+"""
+
+# Execution subsystem
+from .execution import PythonExecutor, Timing, TokenUsage
+
+# Tools subsystem
+from .tools import ToolManager, ToolSource, ToolRegistry, ToolSelector, MCPManager
+
+# Workflow subsystem
+from .workflow import WorkflowEngine, StateManager, PlanManager
+
+# Support subsystem
+from .support import ConsoleDisplay, PackageManager
+
+__all__ = [
+    # Execution
+    'PythonExecutor',
+    'Timing',
+    'TokenUsage',
+    # Tools
+    'ToolManager',
+    'ToolSource',
+    'ToolRegistry',
+    'ToolSelector',
+    'MCPManager',
+    # Workflow
+    'WorkflowEngine',
+    'StateManager',
+    'PlanManager',
+    # Support
+    'ConsoleDisplay',
+    'PackageManager'
+]

managers/execution/__init__.py

@@ -0,0 +1,8 @@
+"""
+Execution subsystem - Python execution and monitoring.
+"""
+
+from .python_executor import PythonExecutor
+from .monitoring import Timing, TokenUsage
+
+__all__ = ['PythonExecutor', 'Timing', 'TokenUsage']

managers/execution/monitoring.py

@@ -0,0 +1,52 @@
+"""
+Monitoring utilities for tracking execution timing and token usage.
+"""
+
+from dataclasses import dataclass, field
+
+__all__ = ["Timing", "TokenUsage"]
+
+
+@dataclass
+class TokenUsage:
+    """
+    Contains the token usage information for a given step or run.
+    """
+
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int = field(init=False)
+
+    def __post_init__(self):
+        self.total_tokens = self.input_tokens + self.output_tokens
+
+    def dict(self):
+        return {
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "total_tokens": self.total_tokens,
+        }
+
+
+@dataclass
+class Timing:
+    """
+    Contains the timing information for a given step or run.
+    """
+
+    start_time: float
+    end_time: float | None = None
+
+    @property
+    def duration(self):
+        return None if self.end_time is None else self.end_time - self.start_time
+
+    def dict(self):
+        return {
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration": self.duration,
+        }
+
+    def __repr__(self) -> str:
+        return f"Timing(start_time={self.start_time}, end_time={self.end_time}, duration={self.duration})"

managers/execution/python_executor.py

@@ -0,0 +1,138 @@
+"""
+Python Executor - Persistent Python execution environment.
+Similar to Jupyter kernel or smolagents LocalPythonExecutor.
+"""
+
+import io
+import sys
+from typing import Any, Dict
+
+__all__ = ["PythonExecutor"]
+
+
+class PythonExecutor:
+    """
+    Python execution environment similar to smolagents LocalPythonExecutor.
+    Maintains persistent state and tool namespace across executions.
+    """
+
+    def __init__(self):
+        self.namespace = {}
+        self.reset_environment()
+
+    def reset_environment(self):
+        """Reset the execution environment."""
+        # Start with basic Python builtins
+        self.namespace = {
+            "__builtins__": __builtins__,
+        }
+        # Add print output capture
+        self.namespace["_print_outputs"] = []
+
+        # Override print to capture outputs
+        def captured_print(*args, **kwargs):
+            output = " ".join(str(arg) for arg in args)
+            self.namespace["_print_outputs"].append(output)
+            print(*args, **kwargs)  # Also print to console
+
+        self.namespace["print"] = captured_print
+
+        # Pre-import commonly used libraries
+        try:
+            import pandas as pd
+            import numpy as np
+            import os
+            from pathlib import Path
+            self.namespace["pd"] = pd
+            self.namespace["np"] = np
+            self.namespace["os"] = os
+            self.namespace["Path"] = Path
+
+            # Add file I/O helper functions
+            def write_text_file(filepath, content):
+                """Write text content to a file."""
+                with open(filepath, 'w') as f:
+                    f.write(content)
+                return f"File written to {filepath}"
+
+            def write_dataframe_to_csv(df, filepath, **kwargs):
+                """Write pandas DataFrame to CSV file."""
+                df.to_csv(filepath, **kwargs)
+                return f"DataFrame written to {filepath}"
+
+            def write_dataframe_to_tsv(df, filepath, **kwargs):
+                """Write pandas DataFrame to TSV file."""
+                df.to_csv(filepath, sep='\t', **kwargs)
+                return f"DataFrame written to {filepath}"
+
+            def create_directory(dirpath):
+                """Create directory if it doesn't exist."""
+                Path(dirpath).mkdir(parents=True, exist_ok=True)
+                return f"Directory created: {dirpath}"
+
+            self.namespace["write_text_file"] = write_text_file
+            self.namespace["write_dataframe_to_csv"] = write_dataframe_to_csv
+            self.namespace["write_dataframe_to_tsv"] = write_dataframe_to_tsv
+            self.namespace["create_directory"] = create_directory
+
+        except ImportError as e:
+            print(f"Warning: Could not import library: {e}")
+
+    def send_functions(self, functions: Dict[str, Any]):
+        """Inject functions into the execution namespace."""
+        for func_name, func in functions.items():
+            self.namespace[func_name] = func
+
+    def send_tools(self, tools: Dict[str, Any]):
+        """Legacy method for backward compatibility - redirects to send_functions."""
+        self.send_functions(tools)
+
+    def send_variables(self, variables: Dict[str, Any]):
+        """Inject variables into the execution namespace."""
+        if variables:
+            print(f"📝 Injecting {len(variables)} variables into Python executor...")
+            self.namespace.update(variables)
+
+    def __call__(self, code: str) -> Any:
+        """Execute code in the persistent namespace."""
+        return self.execute(code)
+
+    def execute(self, code: str) -> Any:
+        """
+        Execute Python code in the persistent environment.
+        Similar to smolagents python execution.
+        """
+        try:
+            # Clear previous print outputs
+            self.namespace["_print_outputs"] = []
+
+            # Capture stdout
+            old_stdout = sys.stdout
+            sys.stdout = captured_output = io.StringIO()
+
+            try:
+                # Execute the code in our persistent namespace
+                exec(code, self.namespace)
+
+                # Get any stdout output
+                stdout_output = captured_output.getvalue()
+
+                # Get print outputs from our custom print function
+                print_outputs = self.namespace.get("_print_outputs", [])
+
+                # Combine outputs
+                all_outputs = []
+                if stdout_output.strip():
+                    all_outputs.append(stdout_output.strip())
+                if print_outputs:
+                    all_outputs.extend(print_outputs)
+
+                result = "\n".join(all_outputs) if all_outputs else "Code executed successfully"
+
+                return result
+
+            finally:
+                sys.stdout = old_stdout
+
+        except Exception as e:
+            return f"Error: {str(e)}"

managers/support/__init__.py

@@ -0,0 +1,8 @@
+"""
+Support subsystem - Console display and package management.
+"""
+
+from .console_display import ConsoleDisplay
+from .package_manager import PackageManager
+
+__all__ = ['ConsoleDisplay', 'PackageManager']

managers/support/console_display.py

@@ -0,0 +1,140 @@
+"""
+Console Display Manager for CodeAct Agent.
+Handles all console output and display formatting.
+"""
+
+from typing import List, Dict
+from rich import box
+from rich.console import Console
+from rich.panel import Panel
+from rich.rule import Rule
+from rich.syntax import Syntax
+from rich.table import Table
+from rich.text import Text
+from ..workflow.plan_manager import PlanManager
+
+
+class ConsoleDisplay:
+    """Handles all console output and display formatting."""
+
+    def __init__(self):
+        self.console = Console()
+
+    def print_task_header(self, query: str):
+        """Print the task header panel."""
+        self.console.print(
+            Panel(
+                f"\n[bold]{query}\n",
+                title="[bold]Code Agent - New Run",
+                subtitle="Starting agent execution...",
+                border_style="yellow",
+                subtitle_align="left",
+            )
+        )
+
+    def print_tools_info(self, total_tools: int, regular_tools: List[str], mcp_tools: Dict):
+        """Print information about loaded tools."""
+        self.console.print(f"🛠️  Loaded {total_tools} total functions:")
+        if regular_tools:
+            self.console.print(f"   📋 Regular tools ({len(regular_tools)}): {regular_tools}")
+
+        if mcp_tools:
+            self.console.print(f"   🔗 MCP tools ({len(mcp_tools)}):")
+            for tool_name, tool_info in mcp_tools.items():
+                server_name = tool_info.get('server', 'unknown')
+                self.console.print(f"      • {tool_name} (from {server_name} server)")
+
+    def print_packages_info(self, imported_packages: List[str], failed_packages: List[tuple]):
+        """Print information about imported packages."""
+        if imported_packages:
+            self.console.print(f"📦 Imported {len(imported_packages)} packages: {imported_packages}")
+
+        for package_name, error in failed_packages:
+            self.console.print(f"⚠️  Failed to import {package_name}: {error}")
+
+    def print_reasoning(self, content: str, step_count: int):
+        """Print agent reasoning panel."""
+        self.console.print(
+            Panel(
+                Text(content, style="italic"),
+                title=f"[bold]🧠 Agent Reasoning - Step {step_count}",
+                border_style="yellow",
+                box=box.SIMPLE,
+            )
+        )
+
+    def print_plan(self, plan: str):
+        """Print current plan panel with progress statistics."""
+        # Get plan progress statistics
+        progress = PlanManager.get_plan_progress(plan)
+        progress_text = f"({progress['completed']}/{progress['total']} completed)"
+
+        self.console.print(
+            Panel(
+                Text(plan, style="cyan"),
+                title=f"[bold]📋 Current Plan {progress_text}",
+                border_style="cyan",
+                box=box.SIMPLE,
+            )
+        )
+
+    def print_code_execution(self, code: str, step_count: int):
+        """Print code execution panel."""
+        self.console.print(f"⚡ Step {step_count}: Executing code...")
+        self.console.print(
+            Panel(
+                Syntax(code, "python", theme="monokai", word_wrap=True),
+                title=f"[bold]💻 Code Execution - Step {step_count}",
+                title_align="left",
+                box=box.HORIZONTALS,
+            )
+        )
+
+    def print_solution(self, solution: str, step_count: int):
+        """Print final solution panel."""
+        self.console.print(f"🎯 Step {step_count}: Providing final solution...")
+        self.console.print(
+            Panel(
+                Text(solution, style="bold green"),
+                title=f"[bold]✅ Final Solution - Step {step_count}",
+                border_style="green",
+            )
+        )
+
+    def print_error(self, error: str, step_count: int):
+        """Print error panel."""
+        self.console.print(f"❌ Step {step_count}: Error encountered...")
+        self.console.print(
+            Panel(
+                Text(error, style="bold red"),
+                title=f"[bold]⚠️ Error - Step {step_count}",
+                border_style="red",
+            )
+        )
+
+    def print_execution_result(self, result: str, step_count: int):
+        """Print execution result panel."""
+        self.console.print(
+            Panel(
+                Syntax(result, "text", theme="github-dark", word_wrap=True),
+                title=f"[bold]📊 Execution Result - Step {step_count}",
+                border_style="cyan",
+                title_align="left",
+            )
+        )
+
+    def print_execution_summary(self, step_count: int, error_count: int, duration: float):
+        """Print execution summary table."""
+        self.console.print(Rule(title="Execution Summary", style="yellow"))
+
+        summary_table = Table(show_header=False, box=box.SIMPLE)
+        summary_table.add_column("Metric", style="cyan")
+        summary_table.add_column("Value", style="bold")
+
+        summary_table.add_row("✅ Total steps", str(step_count))
+        summary_table.add_row("⚠️  Total errors", str(error_count))
+        summary_table.add_row("⏱️  Total execution time", f"{duration:.2f} seconds")
+        if step_count > 0:
+            summary_table.add_row("📊 Average time per step", f"{duration/step_count:.2f}s")
+
+        self.console.print(summary_table)

managers/support/package_manager.py

@@ -0,0 +1,58 @@
+"""
+Package Manager for CodeAct Agent.
+Handles Python package management and imports.
+"""
+
+from typing import Dict
+from core.constants import LIBRARY_CONTENT_DICT
+
+
+class PackageManager:
+    """Manages Python packages available to the agent."""
+
+    def __init__(self, default_packages: Dict[str, str] = None):
+        self.default_packages = default_packages or LIBRARY_CONTENT_DICT.copy()
+        self.custom_packages = {}
+
+    def add_packages(self, packages: Dict[str, str]) -> bool:
+        """Add new packages to the available packages."""
+        try:
+            if not isinstance(packages, dict):
+                raise ValueError("Packages must be a dictionary with package name as key and description as value")
+
+            for package_name, description in packages.items():
+                if not isinstance(package_name, str) or not isinstance(description, str):
+                    print("Warning: Skipping invalid package entry - package_name and description must be strings")
+                    continue
+
+                self.custom_packages[package_name] = description
+                print(f"Added package '{package_name}': {description}")
+
+            print(f"Successfully added {len(packages)} package(s) to the library")
+            return True
+
+        except Exception as e:
+            print(f"Error adding packages: {e}")
+            return False
+
+    def get_all_packages(self) -> Dict[str, str]:
+        """Get all available packages (default + custom)."""
+        all_packages = self.default_packages.copy()
+        all_packages.update(self.custom_packages)
+        return all_packages
+
+    def import_packages(self, python_executor) -> tuple:
+        """Import all packages into the executor environment."""
+        all_packages = self.get_all_packages()
+        imported_packages = []
+        failed_packages = []
+
+        for package_name in all_packages.keys():
+            try:
+                import_code = f"import {package_name}"
+                python_executor(import_code)
+                imported_packages.append(package_name)
+            except Exception as e:
+                failed_packages.append((package_name, str(e)))
+
+        return imported_packages, failed_packages

managers/tools/__init__.py

@@ -0,0 +1,23 @@
+"""
+Tools subsystem - Tool management, registry, and integrations.
+"""
+
+from .tool_manager import ToolManager, ToolSource, ToolInfo
+from .tool_registry import ToolRegistry, tool, discover_tools_in_module, create_module2api_from_functions
+from .tool_selector import ToolSelector
+from .mcp_manager import MCPManager
+from .builtin_tools import FUNCTION_REGISTRY, get_all_tool_functions
+
+__all__ = [
+    'ToolManager',
+    'ToolSource',
+    'ToolInfo',
+    'ToolRegistry',
+    'tool',
+    'discover_tools_in_module',
+    'create_module2api_from_functions',
+    'ToolSelector',
+    'MCPManager',
+    'FUNCTION_REGISTRY',
+    'get_all_tool_functions'
+]

managers/tools/builtin_tools.py

@@ -0,0 +1,34 @@
+"""
+Built-in tool functions for CodeAct agent.
+Example tools that demonstrate the tool system.
+"""
+
+from typing import Any, List
+
+__all__ = ["FUNCTION_REGISTRY", "get_all_tool_functions",
+           "add_numbers", "multiply_numbers"]
+
+
+def add_numbers(a: int, b: int) -> int:
+    """Add two numbers together."""
+    return a + b
+
+
+def multiply_numbers(a: int, b: int) -> int:
+    """Multiply two numbers together."""
+    return a * b
+
+
+# ====================
+# FUNCTION REGISTRY
+# ====================
+
+FUNCTION_REGISTRY = {
+    "add_numbers": add_numbers,
+    "multiply_numbers": multiply_numbers,
+}
+
+
+def get_all_tool_functions() -> List[Any]:
+    """Get all functions from FUNCTION_REGISTRY."""
+    return list(FUNCTION_REGISTRY.values())

managers/tools/mcp_manager.py

@@ -0,0 +1,588 @@
+"""
+MCP Manager for CodeAct Agent.
+Manages MCP (Model Context Protocol) tools and servers.
+"""
+
+import os
+import sys
+import types
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+from rich.console import Console
+
+
+class MCPManager:
+    """Manages MCP (Model Context Protocol) tools and servers."""
+
+    def __init__(self, console_display=None):
+        self.mcp_functions = {}
+        self.console = console_display.console if console_display else Console()
+
+    def has_mcp_functions(self) -> bool:
+        """Check if MCP functions are available."""
+        return bool(self.mcp_functions)
+
+    def group_tools_by_server(self, mcp_tools: Dict[str, dict]) -> Dict[str, List[Dict]]:
+        """Group MCP tools by server name."""
+        servers = {}
+        for tool_name, tool_info in mcp_tools.items():
+            server_name = tool_info.get('server', 'unknown')
+            if server_name not in servers:
+                servers[server_name] = []
+            servers[server_name].append({
+                'name': tool_name,
+                'description': tool_info.get('description', 'No description')
+            })
+        return servers
+
+    def add_mcp(self, config_path: str = "./mcp_config.yaml", tool_registry=None) -> None:
+        """Add MCP tools from configuration file."""
+        try:
+            import asyncio
+            import yaml
+        except ImportError as e:
+            raise ImportError(f"Required packages not available: {e}. Install with: pip install pyyaml") from e
+
+        try:
+            import nest_asyncio
+            from mcp import ClientSession
+            from mcp.client.stdio import StdioServerParameters, stdio_client
+            from mcp.client.streamable_http import streamablehttp_client
+            from langchain_mcp_adapters.tools import _list_all_tools
+            nest_asyncio.apply()
+        except ImportError as e:
+            raise ImportError(f"MCP packages not available: {e}. Install with: pip install mcp langchain-mcp-adapters") from e
+
+        def discover_mcp_tools_sync(server_params: StdioServerParameters) -> List[dict]:
+            """Discover available tools from MCP server synchronously."""
+            try:
+                async def _discover_async():
+                    async with stdio_client(server_params) as (reader, writer):
+                        async with ClientSession(reader, writer) as session:
+                            await session.initialize()
+
+                            tools_result = await session.list_tools()
+                            tools = tools_result.tools if hasattr(tools_result, "tools") else tools_result
+                            print(tools)
+
+                            discovered_tools = []
+                            for tool in tools:
+                                if hasattr(tool, "name"):
+                                    # Ensure description is never empty or None
+                                    description = getattr(tool, 'description', None)
+                                    if not description or description.strip() == "":
+                                        # Generate description from tool name
+                                        formatted_name = tool.name.replace('_', ' ').title()
+                                        description = f"MCP tool: {formatted_name}"
+
+                                    discovered_tools.append({
+                                        "name": tool.name,
+                                        "description": description,
+                                        "inputSchema": tool.inputSchema,
+                                    })
+                                else:
+                                    print(f"Warning: Skipping tool with no name attribute: {tool}")
+
+                            return discovered_tools
+
+                return asyncio.run(_discover_async())
+            except Exception as e:
+                print(f"Failed to discover tools: {e}")
+                return []
+
+        def discover_remote_mcp_tools_sync(url: str) -> List[dict]:
+            """Discover available tools from remote MCP server synchronously."""
+            try:
+                async def _discover_remote_async():
+                    async with streamablehttp_client(url) as (read, write, _):
+                        async with ClientSession(read, write) as session:
+                            await session.initialize()
+                            tools = await _list_all_tools(session)
+
+                            discovered_tools = []
+                            for tool in tools:
+                                if hasattr(tool, "name"):
+                                    # Ensure description is never empty or None
+                                    description = getattr(tool, 'description', None)
+                                    if not description or description.strip() == "":
+                                        # Generate description from tool name
+                                        formatted_name = tool.name.replace('_', ' ').title()
+                                        description = f"MCP tool: {formatted_name}"
+
+                                    discovered_tools.append({
+                                        "name": tool.name,
+                                        "description": description,
+                                        "inputSchema": tool.inputSchema,
+                                    })
+                                else:
+                                    print(f"Warning: Skipping tool with no name attribute: {tool}")
+
+                            return discovered_tools
+
+                return asyncio.run(_discover_remote_async())
+            except Exception as e:
+                print(f"Failed to discover remote tools from {url}: {e}")
+                return []
+
+        def make_mcp_wrapper(cmd: str, args: List[str], tool_name: str, doc: str, env_vars: dict = None):
+            """Create a synchronous wrapper for an async MCP tool call."""
+
+            def sync_tool_wrapper(**kwargs):
+                """Synchronous wrapper for MCP tool execution."""
+                try:
+                    server_params = StdioServerParameters(command=cmd, args=args, env=env_vars)
+
+                    async def async_tool_call():
+                        async with stdio_client(server_params) as (reader, writer):
+                            async with ClientSession(reader, writer) as session:
+                                await session.initialize()
+                                result = await session.call_tool(tool_name, kwargs)
+                                content = result.content[0]
+                                if hasattr(content, "model_dump_json"):
+                                    return content.model_dump_json()
+                                elif hasattr(content, "json"):
+                                    return content.json()
+                                return content.text
+
+                    try:
+                        loop = asyncio.get_running_loop()
+                        return loop.create_task(async_tool_call())
+                    except RuntimeError:
+                        return asyncio.run(async_tool_call())
+
+                except Exception as e:
+                    raise RuntimeError(f"MCP tool execution failed for '{tool_name}': {e}") from e
+
+            sync_tool_wrapper.__name__ = tool_name
+            sync_tool_wrapper.__doc__ = doc
+            return sync_tool_wrapper
+
+        def make_remote_mcp_wrapper(url: str, tool_name: str, doc: str):
+            """Create a synchronous wrapper for an async remote MCP tool call."""
+
+            def sync_tool_wrapper(**kwargs):
+                """Synchronous wrapper for remote MCP tool execution."""
+                try:
+                    async def async_tool_call():
+                        async with streamablehttp_client(url) as (read, write, _):
+                            async with ClientSession(read, write) as session:
+                                await session.initialize()
+                                result = await session.call_tool(tool_name, kwargs)
+                                content = result.content[0]
+                                if hasattr(content, "model_dump_json"):
+                                    return content.model_dump_json()
+                                elif hasattr(content, "json"):
+                                    return content.json()
+                                return content.text
+
+                    try:
+                        loop = asyncio.get_running_loop()
+                        return loop.create_task(async_tool_call())
+                    except RuntimeError:
+                        return asyncio.run(async_tool_call())
+
+                except Exception as e:
+                    raise RuntimeError(f"Remote MCP tool execution failed for '{tool_name}': {e}") from e
+
+            sync_tool_wrapper.__name__ = tool_name
+            sync_tool_wrapper.__doc__ = doc
+            return sync_tool_wrapper
+
+        # Load and validate configuration
+        try:
+            config_content = Path(config_path).read_text(encoding="utf-8")
+            cfg = yaml.safe_load(config_content) or {}
+        except FileNotFoundError:
+            raise FileNotFoundError(f"MCP config file not found: {config_path}") from None
+        except yaml.YAMLError as e:
+            raise yaml.YAMLError(f"Invalid YAML in MCP config: {e}") from e
+
+        mcp_servers = cfg.get("mcp_servers", {})
+        if not mcp_servers:
+            print("Warning: No MCP servers found in configuration")
+            return
+
+        # Process each MCP server configuration
+        for server_name, server_meta in mcp_servers.items():
+            if not server_meta.get("enabled", True):
+                continue
+
+            # Check if this is a remote server configuration
+            remote_url = server_meta.get("url")
+            if remote_url:
+                # Handle remote MCP server
+                self._process_remote_server(server_name, server_meta, remote_url, tool_registry, discover_remote_mcp_tools_sync, make_remote_mcp_wrapper)
+                continue
+
+            # Handle local MCP server (existing logic)
+            # Validate command configuration
+            cmd_list = server_meta.get("command", [])
+            if not cmd_list or not isinstance(cmd_list, list):
+                print(f"Warning: Invalid command configuration for server '{server_name}'")
+                continue
+
+            cmd, *args = cmd_list
+
+            # Process environment variables
+            env_vars = server_meta.get("env", {})
+            if env_vars:
+                processed_env = {}
+                for key, value in env_vars.items():
+                    if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
+                        var_name = value[2:-1]
+                        processed_env[key] = os.getenv(var_name, "")
+                    else:
+                        processed_env[key] = value
+                env_vars = processed_env
+
+            # Create module namespace for this MCP server
+            mcp_module_name = f"mcp_servers.{server_name}"
+            if mcp_module_name not in sys.modules:
+                sys.modules[mcp_module_name] = types.ModuleType(mcp_module_name)
+            server_module = sys.modules[mcp_module_name]
+
+            tools_config = server_meta.get("tools", [])
+
+            # Auto-discover tools if not manually configured
+            if not tools_config:
+                try:
+                    server_params = StdioServerParameters(command=cmd, args=args, env=env_vars)
+                    tools_config = discover_mcp_tools_sync(server_params)
+
+                    if tools_config:
+                        print(f"🔍 Discovered {len(tools_config)} tools from {server_name} MCP server")
+                    else:
+                        print(f"Warning: No tools discovered from {server_name} MCP server")
+                        continue
+
+                except Exception as e:
+                    print(f"Failed to discover tools for {server_name}: {e}")
+                    continue
+
+            # Register each tool
+            tools_added = 0
+            for tool_meta in tools_config:
+                if isinstance(tool_meta, dict) and "biomni_name" in tool_meta:
+                    # Manual tool definition (Biomni-style)
+                    tool_name = tool_meta.get("biomni_name")
+                    description = tool_meta.get("description", f"MCP tool: {tool_name}")
+                    parameters = tool_meta.get("parameters", {})
+                    required_param_names = []
+                    for param_name, param_spec in parameters.items():
+                        if param_spec.get("required", False):
+                            required_param_names.append(param_name)
+                else:
+                    # Auto-discovered tool
+                    tool_name = tool_meta.get("name")
+                    description = tool_meta.get("description", "")
+
+                    # Ensure description is never empty
+                    if not description or description.strip() == "":
+                        formatted_name = tool_name.replace('_', ' ').title()
+                        description = f"MCP tool: {formatted_name}"
+
+                    input_schema = tool_meta.get("inputSchema", {})
+                    parameters = input_schema.get("properties", {})
+                    required_param_names = input_schema.get("required", [])
+
+                if not tool_name:
+                    print(f"Warning: Skipping tool with no name in {server_name}")
+                    continue
+
+                # Create wrapper function
+                wrapper_function = make_mcp_wrapper(cmd, args, tool_name, description, env_vars)
+
+                # Add to module namespace
+                setattr(server_module, tool_name, wrapper_function)
+
+                # Store in MCP functions registry with parameter information
+                self.mcp_functions[tool_name] = {
+                    "function": wrapper_function,
+                    "server": server_name,
+                    "module": mcp_module_name,
+                    "description": description,
+                    "required_parameters": [],  # Will be populated below
+                    "optional_parameters": []   # Will be populated below
+                }
+
+                # Register with tool registry if available
+                if tool_registry:
+                    from .tool_registry import ToolRegistry
+                    # Create tool schema with proper parameter information
+                    required_params = []
+                    optional_params = []
+
+                    for param_name, param_spec in parameters.items():
+                        param_info = {
+                            "name": param_name,
+                            "type": param_spec.get("type", "string"),
+                            "description": param_spec.get("description", f"Parameter {param_name}"),
+                        }
+
+                        # Extract enum/literal values if present
+                        if "enum" in param_spec:
+                            param_info["enum"] = param_spec["enum"]
+
+                        # Handle anyOf schemas (common for optional literal types)
+                        if "anyOf" in param_spec:
+                            # Look for enum in anyOf schemas
+                            for schema_option in param_spec["anyOf"]:
+                                if "enum" in schema_option:
+                                    param_info["enum"] = schema_option["enum"]
+                                    # Update type if specified
+                                    if "type" in schema_option:
+                                        param_info["type"] = schema_option["type"]
+                                    break
+
+                        # Handle oneOf schemas (alternative union syntax)
+                        if "oneOf" in param_spec:
+                            # Look for enum in oneOf schemas
+                            for schema_option in param_spec["oneOf"]:
+                                if "enum" in schema_option:
+                                    param_info["enum"] = schema_option["enum"]
+                                    if "type" in schema_option:
+                                        param_info["type"] = schema_option["type"]
+                                    break
+
+                        # Determine if parameter is required based on:
+                        # 1. Explicit required list (if provided)
+                        # 2. If no default value is present in the schema
+                        is_required = (param_name in required_param_names) or ("default" not in param_spec)
+
+                        if is_required:
+                            required_params.append(param_info)
+                        else:
+                            param_info["default"] = param_spec.get("default")
+                            optional_params.append(param_info)
+
+                    # Create complete tool schema
+                    tool_schema = {
+                        "name": tool_name,
+                        "description": description,
+                        "required_parameters": required_params,
+                        "optional_parameters": optional_params,
+                        "module": mcp_module_name,
+                    }
+
+                    success = tool_registry.register_tool(tool_schema, mcp_module_name)
+                    if success:
+                        tool_registry._name_to_function[tool_name] = wrapper_function
+                        tools_added += 1
+
+                        # Update MCP functions registry with parameter information
+                        self.mcp_functions[tool_name]["required_parameters"] = required_params
+                        self.mcp_functions[tool_name]["optional_parameters"] = optional_params
+
+            if tools_added > 0:
+                print(f"✅ Added {tools_added} MCP tools from {server_name} server")
+
+        print(f"🛠️  Total MCP tools loaded: {len(self.mcp_functions)}")
+
+    def _process_remote_server(self, server_name: str, server_meta: dict, remote_url: str, tool_registry, discover_remote_mcp_tools_sync, make_remote_mcp_wrapper):
+        """Process a remote MCP server configuration."""
+        import sys
+        import types
+
+        # Create module namespace for this remote MCP server
+        mcp_module_name = f"mcp_servers.{server_name}"
+        if mcp_module_name not in sys.modules:
+            sys.modules[mcp_module_name] = types.ModuleType(mcp_module_name)
+        server_module = sys.modules[mcp_module_name]
+
+        tools_config = server_meta.get("tools", [])
+
+        # Auto-discover tools if not manually configured
+        if not tools_config:
+            try:
+                tools_config = discover_remote_mcp_tools_sync(remote_url)
+
+                if tools_config:
+                    print(f"🔍 Discovered {len(tools_config)} tools from {server_name} remote MCP server")
+                else:
+                    print(f"Warning: No tools discovered from {server_name} remote MCP server")
+                    return
+
+            except Exception as e:
+                print(f"Failed to discover tools for remote {server_name}: {e}")
+                return
+
+        # Register each tool
+        tools_added = 0
+        for tool_meta in tools_config:
+            if isinstance(tool_meta, dict) and "biomni_name" in tool_meta:
+                # Manual tool definition (Biomni-style)
+                tool_name = tool_meta.get("biomni_name")
+                description = tool_meta.get("description", f"Remote MCP tool: {tool_name}")
+                parameters = tool_meta.get("parameters", {})
+                required_param_names = []
+                for param_name, param_spec in parameters.items():
+                    if param_spec.get("required", False):
+                        required_param_names.append(param_name)
+            else:
+                # Auto-discovered tool
+                tool_name = tool_meta.get("name")
+                description = tool_meta.get("description", "")
+
+                # Ensure description is never empty
+                if not description or description.strip() == "":
+                    formatted_name = tool_name.replace('_', ' ').title()
+                    description = f"Remote MCP tool: {formatted_name}"
+
+                input_schema = tool_meta.get("inputSchema", {})
+                parameters = input_schema.get("properties", {})
+                required_param_names = input_schema.get("required", [])
+
+            if not tool_name:
+                print(f"Warning: Skipping tool with no name in remote {server_name}")
+                continue
+
+            # Create wrapper function for remote tool
+            wrapper_function = make_remote_mcp_wrapper(remote_url, tool_name, description)
+
+            # Add to module namespace
+            setattr(server_module, tool_name, wrapper_function)
+
+            # Store in MCP functions registry with parameter information
+            self.mcp_functions[tool_name] = {
+                "function": wrapper_function,
+                "server": server_name,
+                "module": mcp_module_name,
+                "description": description,
+                "required_parameters": [],  # Will be populated below
+                "optional_parameters": [],  # Will be populated below
+                "remote_url": remote_url
+            }
+
+            # Register with tool registry if available
+            if tool_registry:
+                from .tool_registry import ToolRegistry
+                # Create tool schema with proper parameter information
+                required_params = []
+                optional_params = []
+
+                for param_name, param_spec in parameters.items():
+                    param_info = {
+                        "name": param_name,
+                        "type": param_spec.get("type", "string"),
+                        "description": param_spec.get("description", f"Parameter {param_name}"),
+                    }
+
+                    # Extract enum/literal values if present
+                    if "enum" in param_spec:
+                        param_info["enum"] = param_spec["enum"]
+
+                    # Handle anyOf schemas (common for optional literal types)
+                    if "anyOf" in param_spec:
+                        # Look for enum in anyOf schemas
+                        for schema_option in param_spec["anyOf"]:
+                            if "enum" in schema_option:
+                                param_info["enum"] = schema_option["enum"]
+                                # Update type if specified
+                                if "type" in schema_option:
+                                    param_info["type"] = schema_option["type"]
+                                break
+
+                    # Handle oneOf schemas (alternative union syntax)
+                    if "oneOf" in param_spec:
+                        # Look for enum in oneOf schemas
+                        for schema_option in param_spec["oneOf"]:
+                            if "enum" in schema_option:
+                                param_info["enum"] = schema_option["enum"]
+                                if "type" in schema_option:
+                                    param_info["type"] = schema_option["type"]
+                                break
+
+                    # Determine if parameter is required based on:
+                    # 1. Explicit required list (if provided)
+                    # 2. If no default value is present in the schema
+                    is_required = (param_name in required_param_names) or ("default" not in param_spec)
+
+                    if is_required:
+                        required_params.append(param_info)
+                    else:
+                        param_info["default"] = param_spec.get("default")
+                        optional_params.append(param_info)
+
+                # Create complete tool schema
+                tool_schema = {
+                    "name": tool_name,
+                    "description": description,
+                    "required_parameters": required_params,
+                    "optional_parameters": optional_params,
+                    "module": mcp_module_name,
+                }
+
+                success = tool_registry.register_tool(tool_schema, mcp_module_name)
+                if success:
+                    tool_registry._name_to_function[tool_name] = wrapper_function
+                    tools_added += 1
+
+                    # Update MCP functions registry with parameter information
+                    self.mcp_functions[tool_name]["required_parameters"] = required_params
+                    self.mcp_functions[tool_name]["optional_parameters"] = optional_params
+
+        if tools_added > 0:
+            print(f"✅ Added {tools_added} remote MCP tools from {server_name} server")
+
+    def list_mcp_tools(self) -> Dict[str, dict]:
+        """List all loaded MCP tools."""
+        return self.mcp_functions.copy()
+
+    def remove_mcp_tool(self, tool_name: str, tool_registry=None) -> bool:
+        """Remove an MCP tool by name."""
+        if not self.has_mcp_functions() or tool_name not in self.mcp_functions:
+            return False
+
+        # Remove from tool registry
+        if tool_registry:
+            tool_registry.remove_tool_by_name(tool_name)
+
+        # Remove from MCP functions
+        del self.mcp_functions[tool_name]
+        return True
+
+    def show_mcp_status(self) -> None:
+        """Display detailed MCP status information to the user."""
+        if not self.has_mcp_functions():
+            self.console.print("🔗 No MCP tools loaded")
+            return
+
+        mcp_tools = self.mcp_functions
+        if not mcp_tools:
+            self.console.print("🔗 MCP system initialized but no tools loaded")
+            return
+
+        # Group tools by server
+        servers = self.group_tools_by_server(mcp_tools)
+
+        # Display server information
+        self.console.print(f"\n🔗 MCP Status Report:")
+        self.console.print(f"   📊 Total servers: {len(servers)}")
+        self.console.print(f"   🛠️  Total MCP tools: {len(mcp_tools)}")
+
+        for server_name, tools in servers.items():
+            self.console.print(f"\n   📡 Server: {server_name}")
+            self.console.print(f"      Status: ✅ Active ({len(tools)} tools)")
+            for tool in tools:
+                self.console.print(f"      • {tool['name']}: {tool['description']}")
+
+    def get_mcp_summary(self) -> Dict[str, any]:
+        """Get a summary of MCP tools for programmatic access."""
+        if not self.has_mcp_functions():
+            return {"total_tools": 0, "servers": {}, "tools": {}}
+
+        mcp_tools = self.mcp_functions
+        # Group tools by server but only get tool names
+        servers = {}
+        for tool_name, tool_info in mcp_tools.items():
+            server_name = tool_info.get('server', 'unknown')
+            if server_name not in servers:
+                servers[server_name] = []
+            servers[server_name].append(tool_name)
+
+        return {
+            "total_tools": len(mcp_tools),
+            "total_servers": len(servers),
+            "servers": servers,
+            "tools": {name: info.get('description', '') for name, info in mcp_tools.items()}
+        }

managers/tools/tool_manager.py

@@ -0,0 +1,521 @@
+"""
+Tool Manager for CodeAct Agent.
+Unified management system for all types of tools: local functions, decorated tools, and MCP tools.
+"""
+
+from typing import Dict, List, Optional, Callable, Any, Union
+from enum import Enum
+from dataclasses import dataclass
+import re
+
+# Import components
+from .tool_registry import ToolRegistry, create_module2api_from_functions
+from .mcp_manager import MCPManager
+
+
+class ToolSource(Enum):
+    """Enumeration of tool sources."""
+    LOCAL = "local"
+    DECORATED = "decorated"
+    MCP = "mcp"
+    ALL = "all"
+
+
+@dataclass
+class ToolInfo:
+    """Comprehensive tool information."""
+    name: str
+    description: str
+    source: ToolSource
+    function: Optional[Callable] = None
+    schema: Optional[Dict] = None
+    server: Optional[str] = None  # For MCP tools
+    module: Optional[str] = None
+    required_parameters: List[Dict] = None
+    optional_parameters: List[Dict] = None
+
+    def __post_init__(self):
+        if self.required_parameters is None:
+            self.required_parameters = []
+        if self.optional_parameters is None:
+            self.optional_parameters = []
+
+
+class ToolManager:
+    """
+    Unified tool management system for CodeAct Agent.
+
+    Manages all types of tools:
+    - Local functions (legacy function registry)
+    - Decorated tools (@tool decorator)
+    - MCP tools (Model Context Protocol)
+
+    Provides a single, consistent interface for tool operations.
+    """
+
+    def __init__(self, console_display=None):
+        """
+        Initialize the ToolManager.
+
+        Args:
+            console_display: Optional console display for MCP status output
+        """
+        # Core components
+        self.tool_registry = ToolRegistry()
+        self.mcp_manager = MCPManager(console_display)
+
+        # Unified tool catalog
+        self._tool_catalog: Dict[str, ToolInfo] = {}
+
+        # Legacy function registry (for backward compatibility)
+        self._legacy_functions: Dict[str, Callable] = {}
+
+        # Initialize with decorated tools from function_tools.py
+        self._discover_decorated_tools()
+
+    # ====================
+    # CORE TOOL MANAGEMENT
+    # ====================
+
+    def add_tool(self, tool: Union[Callable, Dict], name: str = None,
+                 description: str = None, source: ToolSource = ToolSource.LOCAL) -> bool:
+        """
+        Add a tool to the manager.
+
+        Args:
+            tool: Either a callable function or a tool schema dict
+            name: Optional custom name (defaults to function.__name__)
+            description: Optional description (defaults to function.__doc__)
+            source: Source type of the tool
+
+        Returns:
+            True if successfully added
+        """
+        try:
+            if callable(tool):
+                # Handle callable functions
+                tool_name = name or tool.__name__
+                tool_desc = description or tool.__doc__ or f"Function {tool.__name__}"
+
+                # Add to tool registry
+                success = self.tool_registry.add_function_directly(tool_name, tool, tool_desc)
+
+                if success:
+                    # Create ToolInfo and add to catalog
+                    tool_info = ToolInfo(
+                        name=tool_name,
+                        description=tool_desc,
+                        source=source,
+                        function=tool,
+                        schema=self._create_schema_from_function(tool_name, tool, tool_desc)
+                    )
+                    self._tool_catalog[tool_name] = tool_info
+
+                    if source == ToolSource.LOCAL:
+                        self._legacy_functions[tool_name] = tool
+
+                    return True
+
+            elif isinstance(tool, dict):
+                # Handle tool schema dictionaries
+                tool_name = tool.get("name") or name
+                tool_desc = tool.get("description") or description
+
+                if not tool_name:
+                    print("Warning: Tool schema must have a name")
+                    return False
+
+                # Register with tool registry
+                success = self.tool_registry.register_tool(tool)
+
+                if success:
+                    tool_info = ToolInfo(
+                        name=tool_name,
+                        description=tool_desc,
+                        source=source,
+                        schema=tool,
+                        required_parameters=tool.get("required_parameters", []),
+                        optional_parameters=tool.get("optional_parameters", [])
+                    )
+                    self._tool_catalog[tool_name] = tool_info
+                    return True
+
+            return False
+
+        except Exception as e:
+            print(f"Error adding tool {name}: {e}")
+            return False
+
+    def remove_tool(self, name: str) -> bool:
+        """Remove a tool by name from all registries."""
+        try:
+            success = False
+
+            # Remove from tool registry
+            if self.tool_registry.remove_tool_by_name(name):
+                success = True
+
+            # Remove from MCP if it's an MCP tool
+            if name in self.mcp_manager.mcp_functions:
+                if self.mcp_manager.remove_mcp_tool(name, self.tool_registry):
+                    success = True
+
+            # Remove from legacy functions
+            if name in self._legacy_functions:
+                del self._legacy_functions[name]
+                success = True
+
+            # Remove from catalog
+            if name in self._tool_catalog:
+                del self._tool_catalog[name]
+                success = True
+
+            return success
+
+        except Exception as e:
+            print(f"Error removing tool {name}: {e}")
+            return False
+
+    def get_tool(self, name: str) -> Optional[ToolInfo]:
+        """Get comprehensive tool information by name."""
+        return self._tool_catalog.get(name)
+
+    def get_tool_function(self, name: str) -> Optional[Callable]:
+        """Get the actual function object by name."""
+        tool_info = self.get_tool(name)
+        if tool_info and tool_info.function:
+            return tool_info.function
+
+        # Check tool registry
+        return self.tool_registry.get_function_by_name(name)
+
+    # ====================
+    # TOOL DISCOVERY AND LISTING
+    # ====================
+
+    def list_tools(self, source: ToolSource = ToolSource.ALL,
+                   include_details: bool = False) -> List[Dict]:
+        """
+        List tools with optional filtering by source.
+
+        Args:
+            source: Filter by tool source (LOCAL, DECORATED, MCP, ALL)
+            include_details: Whether to include detailed information
+
+        Returns:
+            List of tool dictionaries
+        """
+        tools = []
+
+        for tool_name, tool_info in self._tool_catalog.items():
+            if source == ToolSource.ALL or tool_info.source == source:
+                if include_details:
+                    tools.append({
+                        "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_params": len(tool_info.required_parameters),
+                        "optional_params": len(tool_info.optional_parameters)
+                    })
+                else:
+                    tools.append({
+                        "name": tool_info.name,
+                        "description": tool_info.description,
+                        "source": tool_info.source.value
+                    })
+
+        return sorted(tools, key=lambda x: x["name"])
+
+    def search_tools(self, query: str, search_descriptions: bool = True) -> List[Dict]:
+        """
+        Search tools by name and optionally description.
+
+        Args:
+            query: Search query (supports regex)
+            search_descriptions: Whether to also search in descriptions
+
+        Returns:
+            List of matching tools
+        """
+        pattern = re.compile(query, re.IGNORECASE)
+        matching_tools = []
+
+        for tool_name, tool_info in self._tool_catalog.items():
+            match = False
+
+            # Search in name
+            if pattern.search(tool_name):
+                match = True
+
+            # Search in description if enabled
+            elif search_descriptions and pattern.search(tool_info.description or ""):
+                match = True
+
+            if match:
+                matching_tools.append({
+                    "name": tool_info.name,
+                    "description": tool_info.description,
+                    "source": tool_info.source.value,
+                    "server": tool_info.server
+                })
+
+        return sorted(matching_tools, key=lambda x: x["name"])
+
+    def get_tools_by_source(self, source: ToolSource) -> Dict[str, ToolInfo]:
+        """Get all tools from a specific source."""
+        return {
+            name: tool_info
+            for name, tool_info in self._tool_catalog.items()
+            if tool_info.source == source
+        }
+
+    # ====================
+    # MCP INTEGRATION
+    # ====================
+
+    def add_mcp_server(self, config_path: str = "./mcp_config.yaml") -> None:
+        """Add MCP tools from configuration file."""
+        try:
+            # Use MCP manager to load tools
+            self.mcp_manager.add_mcp(config_path, self.tool_registry)
+
+            # Update our catalog with MCP tools
+            mcp_tools = self.mcp_manager.list_mcp_tools()
+            for tool_name, tool_data in mcp_tools.items():
+                tool_info = ToolInfo(
+                    name=tool_name,
+                    description=tool_data.get("description", "MCP tool"),
+                    source=ToolSource.MCP,
+                    function=tool_data.get("function"),
+                    server=tool_data.get("server"),
+                    module=tool_data.get("module"),
+                    required_parameters=tool_data.get("required_parameters", []),
+                    optional_parameters=tool_data.get("optional_parameters", [])
+                )
+                self._tool_catalog[tool_name] = tool_info
+
+        except Exception as e:
+            print(f"Error adding MCP server: {e}")
+
+    def list_mcp_servers(self) -> Dict[str, List[str]]:
+        """List all MCP servers and their tools."""
+        mcp_tools = self.get_tools_by_source(ToolSource.MCP)
+        servers = {}
+
+        for tool_name, tool_info in mcp_tools.items():
+            server_name = tool_info.server or "unknown"
+            if server_name not in servers:
+                servers[server_name] = []
+            servers[server_name].append(tool_name)
+
+        return servers
+
+    def show_mcp_status(self) -> None:
+        """Display detailed MCP status."""
+        self.mcp_manager.show_mcp_status()
+
+    def get_mcp_summary(self) -> Dict[str, Any]:
+        """Get MCP tools summary."""
+        return self.mcp_manager.get_mcp_summary()
+
+    # ====================
+    # TOOL EXECUTION SUPPORT
+    # ====================
+
+    def get_all_functions(self) -> Dict[str, Callable]:
+        """Get all available functions as a dictionary."""
+        functions = {}
+
+        # Add from tool registry
+        functions.update(self.tool_registry.get_all_functions())
+
+        # Add from legacy functions
+        functions.update(self._legacy_functions)
+
+        # Add MCP functions
+        mcp_tools = self.mcp_manager.list_mcp_tools()
+        for tool_name, tool_data in mcp_tools.items():
+            if tool_data.get("function"):
+                functions[tool_name] = tool_data["function"]
+
+        return functions
+
+    def get_tool_schemas(self, openai_format: bool = True) -> List[Dict]:
+        """
+        Get tool schemas for all tools.
+
+        Args:
+            openai_format: Whether to format as OpenAI function schemas
+
+        Returns:
+            List of tool schemas
+        """
+        schemas = []
+
+        for tool_name, tool_info in self._tool_catalog.items():
+            if openai_format:
+                # Convert to OpenAI function schema format
+                schema = {
+                    "type": "function",
+                    "function": {
+                        "name": tool_info.name,
+                        "description": tool_info.description,
+                        "parameters": {
+                            "type": "object",
+                            "properties": {},
+                            "required": []
+                        }
+                    }
+                }
+
+                # Add required parameters
+                for param in tool_info.required_parameters:
+                    param_schema = {
+                        "type": param.get("type", "string"),
+                        "description": param.get("description", "")
+                    }
+                    # Add enum values if present
+                    if "enum" in param:
+                        param_schema["enum"] = param["enum"]
+                    schema["function"]["parameters"]["properties"][param["name"]] = param_schema
+                    schema["function"]["parameters"]["required"].append(param["name"])
+
+                # Add optional parameters
+                for param in tool_info.optional_parameters:
+                    param_schema = {
+                        "type": param.get("type", "string"),
+                        "description": param.get("description", "")
+                    }
+                    # Add enum values if present
+                    if "enum" in param:
+                        param_schema["enum"] = param["enum"]
+                    if "default" in param:
+                        param_schema["default"] = param["default"]
+                    schema["function"]["parameters"]["properties"][param["name"]] = param_schema
+
+                schemas.append(schema)
+            else:
+                # Return raw schema
+                if tool_info.schema:
+                    schemas.append(tool_info.schema)
+
+        return schemas
+
+    # ====================
+    # STATISTICS AND REPORTING
+    # ====================
+
+    def get_tool_statistics(self) -> Dict[str, Any]:
+        """Get comprehensive tool statistics."""
+        stats = {
+            "total_tools": len(self._tool_catalog),
+            "by_source": {source.value: 0 for source in ToolSource if source != ToolSource.ALL},
+            "with_functions": 0,
+            "mcp_servers": len(self.list_mcp_servers()),
+            "tool_registry_size": len(self.tool_registry.tools),
+            "legacy_functions": len(self._legacy_functions)
+        }
+
+        for tool_info in self._tool_catalog.values():
+            stats["by_source"][tool_info.source.value] += 1
+            if tool_info.function:
+                stats["with_functions"] += 1
+
+        return stats
+
+    def validate_tools(self) -> Dict[str, List[str]]:
+        """Validate all tools and return any issues found."""
+        issues = {
+            "missing_functions": [],
+            "missing_descriptions": [],
+            "duplicate_names": [],
+            "invalid_schemas": []
+        }
+
+        seen_names = set()
+        for tool_name, tool_info in self._tool_catalog.items():
+            # Check for duplicates
+            if tool_name in seen_names:
+                issues["duplicate_names"].append(tool_name)
+            seen_names.add(tool_name)
+
+            # Check for missing functions (except MCP tools which may not have direct functions)
+            if not tool_info.function and tool_info.source != ToolSource.MCP:
+                issues["missing_functions"].append(tool_name)
+
+            # Check for missing descriptions
+            if not tool_info.description or tool_info.description.strip() == "":
+                issues["missing_descriptions"].append(tool_name)
+
+        return issues
+
+    # ====================
+    # PRIVATE METHODS
+    # ====================
+
+    def _discover_decorated_tools(self) -> None:
+        """Discover and register tools marked with @tool decorator."""
+        try:
+            from .builtin_tools import get_all_tool_functions
+
+            tool_functions = get_all_tool_functions()
+            for func in tool_functions:
+                name = getattr(func, '_tool_name', func.__name__)
+                description = getattr(func, '_tool_description', func.__doc__ or f"Function {func.__name__}")
+
+                tool_info = ToolInfo(
+                    name=name,
+                    description=description,
+                    source=ToolSource.DECORATED,
+                    function=func,
+                    schema=self._create_schema_from_function(name, func, description)
+                )
+                self._tool_catalog[name] = tool_info
+
+                # Also add to tool registry for consistency
+                self.tool_registry.add_function_directly(name, func, description)
+
+        except ImportError:
+            print("Warning: Could not import builtin_tools module for decorated tool discovery")
+
+    def _create_schema_from_function(self, name: str, function: Callable, description: str) -> Dict:
+        """Create a tool schema from a function object."""
+        return self.tool_registry._create_schema_from_function(name, function, description)
+
+    def _refresh_catalog(self) -> None:
+        """Refresh the tool catalog from all sources."""
+        # Clear current catalog
+        self._tool_catalog.clear()
+
+        # Re-discover decorated tools
+        self._discover_decorated_tools()
+
+        # Re-add MCP tools
+        mcp_tools = self.mcp_manager.list_mcp_tools()
+        for tool_name, tool_data in mcp_tools.items():
+            tool_info = ToolInfo(
+                name=tool_name,
+                description=tool_data.get("description", "MCP tool"),
+                source=ToolSource.MCP,
+                function=tool_data.get("function"),
+                server=tool_data.get("server"),
+                module=tool_data.get("module"),
+                required_parameters=tool_data.get("required_parameters", []),
+                optional_parameters=tool_data.get("optional_parameters", [])
+            )
+            self._tool_catalog[tool_name] = tool_info
+
+    # ====================
+    # LEGACY COMPATIBILITY
+    # ====================
+
+    def add_legacy_functions(self, functions: Dict[str, Callable]) -> int:
+        """Add legacy functions for backward compatibility."""
+        added_count = 0
+        for name, func in functions.items():
+            if self.add_tool(func, name, source=ToolSource.LOCAL):
+                added_count += 1
+        return added_count

managers/tools/tool_registry.py

@@ -0,0 +1,314 @@
+"""
+Tool Registry System for CodeAct - Based on biomni's approach.
+Provides centralized tool management with schema validation and discovery.
+"""
+
+import ast
+import importlib
+import importlib.util
+import inspect
+import os
+import pickle
+from typing import Any, Dict, List, Optional, Callable
+import pandas as pd
+
+__all__ = ["ToolRegistry", "tool", "discover_tools_in_module", "create_module2api_from_functions"]
+
+
+class ToolRegistry:
+    """
+    Central registry for managing tools, similar to biomni's ToolRegistry.
+    Handles tool registration, validation, and lookup by name/ID.
+    """
+
+    def __init__(self, module2api: Optional[Dict[str, List[Dict]]] = None):
+        """
+        Initialize the tool registry.
+
+        Args:
+            module2api: Dictionary mapping module names to lists of tool schemas
+                       e.g., {"module.tools": [{"name": "func1", "description": "...", ...}]}
+        """
+        self.tools = []
+        self.next_id = 0
+        self._name_to_function = {}  # Map tool names to actual functions
+
+        # Register tools from module2api if provided
+        if module2api:
+            for module_name, tool_list in module2api.items():
+                for tool_schema in tool_list:
+                    self.register_tool(tool_schema, module_name)
+
+        # Create document dataframe for retrieval (similar to biomni)
+        self._create_document_df()
+
+    def register_tool(self, tool_schema: Dict, module_name: Optional[str] = None) -> bool:
+        """
+        Register a tool with the registry.
+
+        Args:
+            tool_schema: Tool schema dictionary with name, description, parameters
+            module_name: Optional module name where the tool function is located
+
+        Returns:
+            True if registration successful, False otherwise
+        """
+        if not self.validate_tool(tool_schema):
+            raise ValueError(f"Invalid tool format for {tool_schema.get('name', 'unknown')}")
+
+        # Add unique ID
+        tool_schema = tool_schema.copy()
+        tool_schema["id"] = self.next_id
+        tool_schema["module_name"] = module_name
+
+        # Try to load the actual function if module_name provided
+        if module_name and "name" in tool_schema:
+            try:
+                function = self._load_function_from_module(module_name, tool_schema["name"])
+                if function:
+                    self._name_to_function[tool_schema["name"]] = function
+            except Exception as e:
+                print(f"Warning: Could not load function {tool_schema['name']} from {module_name}: {e}")
+
+        self.tools.append(tool_schema)
+        self.next_id += 1
+        return True
+
+    def validate_tool(self, tool_schema: Dict) -> bool:
+        """Validate that a tool schema has required fields."""
+        required_keys = ["name", "description"]
+        return all(key in tool_schema for key in required_keys)
+
+    def get_tool_by_name(self, name: str) -> Optional[Dict]:
+        """Get tool schema by name."""
+        for tool in self.tools:
+            if tool["name"] == name:
+                return tool
+        return None
+
+    def get_tool_by_id(self, tool_id: int) -> Optional[Dict]:
+        """Get tool schema by ID."""
+        for tool in self.tools:
+            if tool["id"] == tool_id:
+                return tool
+        return None
+
+    def get_function_by_name(self, name: str) -> Optional[Callable]:
+        """Get the actual function object by name."""
+        return self._name_to_function.get(name)
+
+    def get_all_functions(self) -> Dict[str, Callable]:
+        """Get all registered functions as a dictionary."""
+        return self._name_to_function.copy()
+
+    def list_tools(self) -> List[Dict]:
+        """List all registered tools with basic info."""
+        return [{"name": tool["name"], "id": tool["id"], "description": tool["description"]}
+                for tool in self.tools]
+
+    def list_tool_names(self) -> List[str]:
+        """Get list of all tool names."""
+        return [tool["name"] for tool in self.tools]
+
+    def add_function_directly(self, name: str, function: Callable, description: str = None) -> bool:
+        """
+        Add a function directly to the registry.
+
+        Args:
+            name: Function name
+            function: The callable function
+            description: Optional description, will be extracted from docstring if not provided
+
+        Returns:
+            True if added successfully
+        """
+        if description is None:
+            description = function.__doc__ or f"Function {name}"
+
+        # Create schema from function signature
+        schema = self._create_schema_from_function(name, function, description)
+
+        # Register the tool
+        self.register_tool(schema)
+        self._name_to_function[name] = function
+        return True
+
+    def _create_schema_from_function(self, name: str, function: Callable, description: str) -> Dict:
+        """Create a tool schema from a function object."""
+        sig = inspect.signature(function)
+
+        schema = {
+            "name": name,
+            "description": description,
+            "required_parameters": [],
+            "optional_parameters": []
+        }
+
+        for param_name, param in sig.parameters.items():
+            param_info = {
+                "name": param_name,
+                "type": self._get_param_type(param),
+                "description": f"Parameter {param_name}"
+            }
+
+            if param.default == inspect.Parameter.empty:
+                schema["required_parameters"].append(param_info)
+            else:
+                param_info["default"] = param.default
+                schema["optional_parameters"].append(param_info)
+
+        return schema
+
+    def _get_param_type(self, param: inspect.Parameter) -> str:
+        """Extract parameter type as string."""
+        if param.annotation != inspect.Parameter.empty:
+            if hasattr(param.annotation, '__name__'):
+                return param.annotation.__name__
+            else:
+                return str(param.annotation)
+        return "Any"
+
+    def _load_function_from_module(self, module_name: str, function_name: str) -> Optional[Callable]:
+        """Load a function from a module."""
+        try:
+            module = importlib.import_module(module_name)
+            return getattr(module, function_name, None)
+        except (ImportError, AttributeError):
+            return None
+
+    def _create_document_df(self):
+        """Create a pandas DataFrame for tool retrieval (similar to biomni)."""
+        docs = []
+        for tool in self.tools:
+            doc_content = {
+                "name": tool["name"],
+                "description": tool["description"],
+                "required_parameters": tool.get("required_parameters", []),
+                "optional_parameters": tool.get("optional_parameters", []),
+                "module_name": tool.get("module_name", "")
+            }
+            docs.append([tool["id"], doc_content])
+
+        self.document_df = pd.DataFrame(docs, columns=["docid", "document_content"])
+
+    def remove_tool_by_name(self, name: str) -> bool:
+        """Remove a tool by name."""
+        tool = self.get_tool_by_name(name)
+        if tool:
+            self.tools = [t for t in self.tools if t["name"] != name]
+            self._name_to_function.pop(name, None)
+            self._create_document_df()  # Refresh document df
+            return True
+        return False
+
+    def save_registry(self, filename: str):
+        """Save registry to file."""
+        with open(filename, "wb") as file:
+            pickle.dump(self, file)
+
+    @staticmethod
+    def load_registry(filename: str) -> 'ToolRegistry':
+        """Load registry from file."""
+        with open(filename, "rb") as file:
+            return pickle.load(file)
+
+
+# ====================
+# TOOL DECORATOR
+# ====================
+
+def tool(func: Callable = None, *, name: str = None, description: str = None):
+    """
+    Decorator to mark functions as tools, similar to biomni's @tool decorator.
+
+    Usage:
+        @tool
+        def my_function(x: int) -> int:
+            '''This function does something'''
+            return x * 2
+
+        @tool(name="custom_name", description="Custom description")
+        def another_function():
+            pass
+    """
+    def decorator(f):
+        # Store metadata on the function
+        f._tool_name = name or f.__name__
+        f._tool_description = description or f.__doc__ or f"Function {f.__name__}"
+        f._is_tool = True
+        return f
+
+    if func is None:
+        # Called with arguments: @tool(name="...", description="...")
+        return decorator
+    else:
+        # Called without arguments: @tool
+        return decorator(func)
+
+
+# ====================
+# TOOL DISCOVERY UTILITIES
+# ====================
+
+def discover_tools_in_module(module_path: str) -> List[Callable]:
+    """
+    Discover all functions marked with @tool decorator in a module.
+
+    Args:
+        module_path: Path to the Python module file
+
+    Returns:
+        List of function objects marked as tools
+    """
+    with open(module_path, 'r') as file:
+        tree = ast.parse(file.read(), filename=module_path)
+
+    tool_function_names = []
+
+    # Find functions with @tool decorator
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef):
+            for decorator in node.decorator_list:
+                if (isinstance(decorator, ast.Name) and decorator.id == "tool") or \
+                   (isinstance(decorator, ast.Call) and
+                    isinstance(decorator.func, ast.Name) and decorator.func.id == "tool"):
+                    tool_function_names.append(node.name)
+                    break
+
+    # Import the module and get function objects
+    spec = importlib.util.spec_from_file_location("temp_module", module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    tool_functions = []
+    for name in tool_function_names:
+        func = getattr(module, name, None)
+        if func and hasattr(func, '_is_tool'):
+            tool_functions.append(func)
+
+    return tool_functions
+
+
+def create_module2api_from_functions(functions: List[Callable], module_name: str = "custom_tools") -> Dict[str, List[Dict]]:
+    """
+    Create a module2api dictionary from a list of functions.
+
+    Args:
+        functions: List of function objects
+        module_name: Name to assign to the module
+
+    Returns:
+        Dictionary in module2api format
+    """
+    tool_schemas = []
+
+    for func in functions:
+        name = getattr(func, '_tool_name', func.__name__)
+        description = getattr(func, '_tool_description', func.__doc__ or f"Function {func.__name__}")
+
+        # Create schema from function signature
+        registry = ToolRegistry()
+        schema = registry._create_schema_from_function(name, func, description)
+        tool_schemas.append(schema)
+
+    return {module_name: tool_schemas}

managers/tools/tool_selector.py

@@ -0,0 +1,109 @@
+"""
+Tool Selector for CodeAct Agent.
+Pure LLM-based tool selection mechanism similar to Biomni's prompt_based_retrieval.
+"""
+
+import re
+from typing import Dict, List, Optional
+from langchain_core.messages import HumanMessage
+from langchain_core.language_models.chat_models import BaseChatModel
+
+
+class ToolSelector:
+    """
+    LLM-based tool selection system inspired by Biomni's approach.
+    Uses an LLM to intelligently select the most relevant tools for a given task.
+    """
+
+    def __init__(self, model: BaseChatModel):
+        """
+        Initialize the ToolSelector.
+
+        Args:
+            model: The language model to use for tool selection
+        """
+        self.model = model
+
+    def select_tools_for_task(self, query: str, available_tools: Dict[str, Dict], max_tools: int = 15) -> List[str]:
+        """
+        Use LLM-based selection to choose the most relevant tools for a query.
+        Inspired by Biomni's prompt_based_retrieval mechanism.
+
+        Args:
+            query: The user's query/task description
+            available_tools: Dictionary of {tool_name: tool_info} available
+            max_tools: Maximum number of tools to select
+
+        Returns:
+            List of selected tool names
+        """
+        if not available_tools:
+            return []
+
+        # Format tools for LLM prompt
+        tools_list = self._format_tools_for_prompt(available_tools)
+
+        # Create selection prompt (similar to Biomni's approach)
+        selection_prompt = f"""You are an expert biomedical research assistant. Your task is to select the most relevant tools to help answer a user's query.
+
+USER QUERY: {query}
+
+Below are the available tools. Select items that are directly or indirectly relevant to answering the query.
+Be generous in your selection - include tools that might be useful for the task, even if they're not explicitly mentioned in the query.
+It's better to include slightly more tools than to miss potentially useful ones.
+
+AVAILABLE TOOLS:
+{tools_list}
+
+Select up to {max_tools} tools that would be most helpful for this task.
+
+Respond with ONLY a comma-separated list of the exact tool names, like this:
+tool_name_1, tool_name_2, tool_name_3
+
+Selected tools:"""
+
+        try:
+            # Get LLM response
+            response = self.model.invoke([HumanMessage(content=selection_prompt)])
+            response_content = response.content.strip()
+
+            # Parse the response to extract tool names
+            selected_tools = self._parse_tool_selection_response(response_content, available_tools)
+
+            # Ensure we don't exceed max_tools
+            return selected_tools[:max_tools]
+
+        except Exception as e:
+            print(f"Error in LLM-based tool selection: {e}")
+            # Return all tools if LLM fails (no keyword fallback)
+            return list(available_tools.keys())[:max_tools]
+
+    def _format_tools_for_prompt(self, tools: Dict[str, Dict]) -> str:
+        """Format tools for the LLM prompt."""
+        formatted = []
+        for i, (tool_name, tool_info) in enumerate(tools.items(), 1):
+            description = tool_info.get('description', 'No description available')
+            source = tool_info.get('source', 'unknown')
+            formatted.append(f"{i}. {tool_name} ({source}): {description}")
+        return "\n".join(formatted)
+
+    def _parse_tool_selection_response(self, response: str, available_tools: Dict[str, Dict]) -> List[str]:
+        """Parse the LLM response to extract valid tool names."""
+        selected_tools = []
+
+        # Split by commas and clean up
+        tool_candidates = [name.strip() for name in response.split(',')]
+
+        for candidate in tool_candidates:
+            # Remove any extra characters, numbers, or formatting
+            clean_candidate = re.sub(r'^\d+\.\s*', '', candidate)  # Remove "1. " prefixes
+            clean_candidate = clean_candidate.strip()
+
+            # Check if this matches any available tool (case-insensitive)
+            for tool_name in available_tools.keys():
+                if clean_candidate.lower() == tool_name.lower():
+                    if tool_name not in selected_tools:  # Avoid duplicates
+                        selected_tools.append(tool_name)
+                    break
+
+        return selected_tools

managers/workflow/__init__.py

@@ -0,0 +1,9 @@
+"""
+Workflow subsystem - Workflow execution, state, and plan management.
+"""
+
+from .workflow_engine import WorkflowEngine
+from .state_manager import StateManager
+from .plan_manager import PlanManager
+
+__all__ = ['WorkflowEngine', 'StateManager', 'PlanManager']

managers/workflow/plan_manager.py

@@ -0,0 +1,78 @@
+"""
+Plan Manager for CodeAct Agent.
+Handles plan creation, updates, and progress tracking.
+"""
+
+import re
+from typing import Optional, Dict
+
+
+class PlanManager:
+    """Manages plan creation, updates, and progress tracking."""
+
+    @staticmethod
+    def extract_plan_from_content(content: str) -> Optional[str]:
+        """Extract plan from agent content."""
+        plan_pattern = r'\d+\.\s*\[[^\]]*\]\s*[^\n]+(?:\n\d+\.\s*\[[^\]]*\]\s*[^\n]+)*'
+        matches = re.findall(plan_pattern, content)
+        # Return the last (most recent) plan if multiple plans exist
+        return matches[-1] if matches else None
+
+    @staticmethod
+    def update_plan_for_solution(plan_text: str) -> str:
+        """Update plan to mark all remaining steps as completed when providing final solution."""
+        if not plan_text:
+            return plan_text
+
+        lines = plan_text.split('\n')
+        updated_lines = []
+
+        for line in lines:
+            # Mark any unchecked or failed steps as completed since we're providing final solution
+            if '[ ]' in line or '[✗]' in line:
+                updated_line = re.sub(r'\[\s*[^\]]*\]', '[✓]', line)
+                updated_lines.append(updated_line)
+            else:
+                updated_lines.append(line)
+
+        return '\n'.join(updated_lines)
+
+    @staticmethod
+    def mark_step_completed(plan_text: str, step_description: str) -> str:
+        """Mark a specific step as completed based on its description."""
+        if not plan_text or not step_description:
+            return plan_text
+
+        lines = plan_text.split('\n')
+        updated_lines = []
+
+        for line in lines:
+            # Check if this line contains the step description and is unchecked
+            if step_description.lower() in line.lower() and ('[ ]' in line or '[✗]' in line):
+                updated_line = re.sub(r'\[\s*[^\]]*\]', '[✓]', line)
+                updated_lines.append(updated_line)
+            else:
+                updated_lines.append(line)
+
+        return '\n'.join(updated_lines)
+
+    @staticmethod
+    def get_plan_progress(plan_text: str) -> Dict[str, int]:
+        """Get plan progress statistics."""
+        if not plan_text:
+            return {"total": 0, "completed": 0, "pending": 0, "failed": 0}
+
+        lines = plan_text.split('\n')
+        stats = {"total": 0, "completed": 0, "pending": 0, "failed": 0}
+
+        for line in lines:
+            if re.search(r'\d+\.\s*\[', line):
+                stats["total"] += 1
+                if '[✓]' in line:
+                    stats["completed"] += 1
+                elif '[ ]' in line:
+                    stats["pending"] += 1
+                elif '[✗]' in line:
+                    stats["failed"] += 1
+
+        return stats

managers/workflow/state_manager.py

@@ -0,0 +1,24 @@
+"""
+State Manager for CodeAct Agent.
+Manages agent state creation and manipulation.
+"""
+
+import time
+from typing import List, Dict, Optional
+
+
+class StateManager:
+    """Manages agent state creation and manipulation."""
+
+    @staticmethod
+    def create_state_dict(messages: List = None, step_count: int = 0,
+                         error_count: int = 0, start_time: float = None,
+                         current_plan: str = None) -> Dict:
+        """Create a standardized state dictionary."""
+        return {
+            "messages": messages or [],
+            "step_count": step_count,
+            "error_count": error_count,
+            "start_time": start_time or time.time(),
+            "current_plan": current_plan
+        }

managers/workflow/workflow_engine.py

@@ -0,0 +1,303 @@
+"""
+Workflow Engine Manager for CodeAct Agent.
+Manages the LangGraph workflow execution.
+"""
+
+import re
+import json
+import datetime
+from typing import Dict, Tuple, List, Any
+from pathlib import Path
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.rule import Rule
+from langchain_core.messages import AIMessage, HumanMessage, BaseMessage
+from langgraph.graph import StateGraph, START, END
+from core.types import AgentState, AgentConfig
+from .plan_manager import PlanManager
+
+
+class WorkflowEngine:
+    """Manages the LangGraph workflow execution."""
+
+    def __init__(self, model, config: AgentConfig, console_display, state_manager):
+        self.model = model
+        self.config = config
+        self.console = console_display
+        self.state_manager = state_manager
+        self.plan_manager = PlanManager()
+        self.graph = None
+        self.trace_logs = []  # Store all trace logs
+        self.message_history = []  # Store all messages
+
+    def setup_workflow(self, generate_func, execute_func, should_continue_func):
+        """Setup the LangGraph workflow with provided functions."""
+        workflow = StateGraph(AgentState)
+
+        workflow.add_node("generate", generate_func)
+        workflow.add_node("execute", execute_func)
+
+        workflow.add_edge(START, "generate")
+        workflow.add_edge("execute", "generate")
+
+        workflow.add_conditional_edges("generate", should_continue_func, {
+            "end": END,
+            "execute": "execute"
+        })
+
+        self.graph = workflow.compile()
+
+    def run_workflow(self, initial_state: Dict) -> Tuple:
+        """Execute the workflow and handle display.
+
+        Returns:
+            tuple: (result_content, final_state)
+        """
+        # Clear previous traces for new run
+        self.trace_logs = []
+        self.message_history = []
+
+        # Track if final solution has been provided
+        final_solution_provided = False
+        previous_plan = None
+        displayed_reasoning = set()
+
+        # Stream the workflow execution with monitoring
+        self.console.console.print(Rule(title="Execution Steps", style="yellow"))
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            console=self.console.console,
+            transient=True
+        ) as progress:
+            task = progress.add_task("Executing agent...", total=None)
+
+            final_state = None
+            for s in self.graph.stream(initial_state, stream_mode="values"):
+                step_count = s.get("step_count", 0)
+                current_plan = s.get("current_plan")
+                final_state = s
+
+                progress.update(task, description=f"Step {step_count}")
+
+                message = s["messages"][-1]
+
+                # Serialize and store the message
+                serialized_msg = self._serialize_message(message)
+                self.message_history.append(serialized_msg)
+
+                # Process different types of messages
+                if isinstance(message, AIMessage):
+                    self._process_ai_message(
+                        message, step_count, current_plan, previous_plan,
+                        displayed_reasoning, final_solution_provided
+                    )
+                    if current_plan != previous_plan:
+                        previous_plan = current_plan
+
+                elif "<observation>" in message.content:
+                    self._process_observation_message(message, step_count)
+
+        result_content = final_state["messages"][-1].content if final_state else ""
+        return result_content, final_state
+
+    def _process_ai_message(self, message, step_count, current_plan, previous_plan,
+                           displayed_reasoning, final_solution_provided):
+        """Process AI message and display appropriate panels."""
+        full_content = message.content
+
+        # 1. REASONING: Extract and display agent's thinking
+        thinking_content = self._extract_thinking_content(full_content)
+        if thinking_content and len(thinking_content) > 20:
+            content_hash = hash(thinking_content.strip())
+            if content_hash not in displayed_reasoning:
+                self.console.print_reasoning(thinking_content, step_count)
+                displayed_reasoning.add(content_hash)
+                # Add trace entry for reasoning
+                self._add_trace_entry("reasoning", step_count, thinking_content)
+
+        # 2. PLAN: Show plan only when it has changed
+        if (current_plan and current_plan != previous_plan and
+            self.config.verbose and not final_solution_provided):
+            self.console.print_plan(current_plan)
+            # Add trace entry for plan
+            self._add_trace_entry("plan", step_count, current_plan)
+
+        # 3. ACTION & CODE: Handle different action types
+        if "<execute>" in full_content and "</execute>" in full_content:
+            execute_match = re.search(r"<execute>(.*?)</execute>", full_content, re.DOTALL)
+            if execute_match:
+                code = execute_match.group(1).strip()
+                self.console.print_code_execution(code, step_count)
+                # Add trace entry for code execution
+                self._add_trace_entry("code_execution", step_count, code)
+
+        elif "<solution>" in full_content and "</solution>" in full_content:
+            solution_match = re.search(r"<solution>(.*?)</solution>", full_content, re.DOTALL)
+            if solution_match:
+                # Update plan to mark all remaining steps as completed
+                if current_plan:
+                    updated_plan = self.plan_manager.update_plan_for_solution(current_plan)
+                    if updated_plan != current_plan:
+                        self.console.print_plan(updated_plan)
+
+                solution = solution_match.group(1).strip()
+                self.console.print_solution(solution, step_count)
+                final_solution_provided = True
+                # Add trace entry for solution
+                self._add_trace_entry("solution", step_count, solution)
+
+        elif "<error>" in full_content:
+            error_match = re.search(r"<error>(.*?)</error>", full_content, re.DOTALL)
+            if error_match:
+                error_content = error_match.group(1).strip()
+                self.console.print_error(error_content, step_count)
+                # Add trace entry for error
+                self._add_trace_entry("error", step_count, error_content)
+
+    def _process_observation_message(self, message, step_count):
+        """Process observation message and display results."""
+        obs_match = re.search(r"<observation>(.*?)</observation>", message.content, re.DOTALL)
+        if obs_match:
+            observation = obs_match.group(1).strip()
+            formatted_output = self._truncate_to_20_rows(observation)
+            self.console.print_execution_result(formatted_output, step_count)
+            # Add trace entry for observation
+            self._add_trace_entry("observation", step_count, observation)
+
+    def _extract_thinking_content(self, content: str) -> str:
+        """Extract thinking content from the message, removing tags and plan information."""
+        # Remove specific tags but keep observation content for separate handling
+        content = re.sub(r'</?(execute|solution|error)>', '', content)
+
+        # Remove plan content (numbered lists with checkboxes)
+        plan_pattern = r'\d+\.\s*\[[^\]]*\]\s*[^\n]+(?:\n\d+\.\s*\[[^\]]*\]\s*[^\n]+)*'
+        content = re.sub(plan_pattern, '', content).strip()
+
+        # Remove observation blocks entirely
+        content = re.sub(r'<observation>.*?</observation>', '', content, flags=re.DOTALL)
+
+        # Clean up extra whitespace and empty lines
+        lines = [line.strip() for line in content.split('\n') if line.strip()]
+        return '\n'.join(lines)
+
+    def _truncate_to_20_rows(self, text: str) -> str:
+        """Truncate any text output to show only the first 20 rows."""
+        lines = text.split('\n')
+
+        if len(lines) > 20:
+            truncated = '\n'.join(lines[:20])
+            total_lines = len(lines)
+            truncated += f"\n\n⚠️ Output truncated to 20 rows. Full output contains {total_lines} rows."
+            return truncated
+
+        return text
+
+    def _add_trace_entry(self, step_type: str, step_count: int, content: Any, metadata: Dict = None):
+        """Add an entry to the trace log."""
+        entry = {
+            "timestamp": datetime.datetime.now().isoformat(),
+            "step_count": step_count,
+            "step_type": step_type,
+            "content": content,
+            "metadata": metadata or {}
+        }
+        self.trace_logs.append(entry)
+
+    def _serialize_message(self, message: BaseMessage) -> Dict:
+        """Serialize a message for saving."""
+        if isinstance(message, HumanMessage):
+            msg_type = "human"
+        elif isinstance(message, AIMessage):
+            msg_type = "ai"
+        else:
+            msg_type = "system"
+
+        return {
+            "type": msg_type,
+            "content": message.content,
+            "timestamp": datetime.datetime.now().isoformat()
+        }
+
+    def save_trace_to_file(self, filepath: str = None) -> str:
+        """Save the complete trace log to a JSON file."""
+        if filepath is None:
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            filepath = f"agent_trace_{timestamp}.json"
+
+        trace_data = {
+            "execution_time": datetime.datetime.now().isoformat(),
+            "config": {
+                "max_steps": self.config.max_steps,
+                "timeout_seconds": self.config.timeout_seconds,
+                "verbose": self.config.verbose
+            },
+            "messages": self.message_history,
+            "trace_logs": self.trace_logs
+        }
+
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(trace_data, f, indent=2, ensure_ascii=False)
+
+        return filepath
+
+    def generate_summary(self) -> Dict:
+        """Generate a summary of the agent execution."""
+        summary = {
+            "total_steps": len(self.trace_logs),
+            "message_count": len(self.message_history),
+            "execution_flow": [],
+            "code_executions": [],
+            "observations": [],
+            "errors": [],
+            "final_solution": None
+        }
+
+        for entry in self.trace_logs:
+            step_info = {
+                "step": entry["step_count"],
+                "type": entry["step_type"],
+                "timestamp": entry["timestamp"]
+            }
+
+            if entry["step_type"] == "reasoning":
+                summary["execution_flow"].append({
+                    **step_info,
+                    "reasoning": entry["content"][:200] + "..." if len(entry["content"]) > 200 else entry["content"]
+                })
+            elif entry["step_type"] == "code_execution":
+                summary["code_executions"].append({
+                    **step_info,
+                    "code": entry["content"]
+                })
+            elif entry["step_type"] == "observation":
+                summary["observations"].append({
+                    **step_info,
+                    "output": entry["content"][:500] + "..." if len(entry["content"]) > 500 else entry["content"]
+                })
+            elif entry["step_type"] == "error":
+                summary["errors"].append({
+                    **step_info,
+                    "error": entry["content"]
+                })
+            elif entry["step_type"] == "solution":
+                summary["final_solution"] = {
+                    **step_info,
+                    "solution": entry["content"]
+                }
+
+        return summary
+
+    def save_summary_to_file(self, filepath: str = None) -> str:
+        """Save the execution summary to a JSON file."""
+        if filepath is None:
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            filepath = f"agent_summary_{timestamp}.json"
+
+        summary = self.generate_summary()
+        summary["timestamp"] = datetime.datetime.now().isoformat()
+
+        with open(filepath, 'w', encoding='utf-8') as f:
+            json.dump(summary, f, indent=2, ensure_ascii=False)
+
+        return filepath