agents/code_agent.py

"""
Code Agent - Computational tasks and code execution

The Code Agent is responsible for:
1. Performing mathematical calculations
2. Executing Python code for data analysis
3. Processing numerical data and computations
4. Returning structured computational results
"""

import os
import sys
import io
import contextlib
from typing import Dict, Any, List
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
from langgraph.types import Command
from langchain_groq import ChatGroq
from langchain_core.tools import Tool
from observability import agent_span, tool_span
from dotenv import load_dotenv

# Import calculator tools from the existing tools.py
from tools import get_calculator_tool, get_hub_stats_tool

load_dotenv("env.local")


def create_code_tools() -> List[Tool]:
    """Create LangChain-compatible computational tools"""
    tools = []
    
    # Mathematical calculator tools
    def multiply_func(a: float, b: float) -> str:
        """Multiply two numbers"""
        try:
            result = a * b
            return f"{a} × {b} = {result}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def add_func(a: float, b: float) -> str:
        """Add two numbers"""
        try:
            result = a + b
            return f"{a} + {b} = {result}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def subtract_func(a: float, b: float) -> str:
        """Subtract two numbers"""
        try:
            result = a - b
            return f"{a} - {b} = {result}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def divide_func(a: float, b: float) -> str:
        """Divide two numbers"""
        try:
            if b == 0:
                return "Error: Cannot divide by zero"
            result = a / b
            return f"{a} ÷ {b} = {result}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def modulus_func(a: int, b: int) -> str:
        """Get the modulus of two numbers"""
        try:
            if b == 0:
                return "Error: Cannot modulo by zero"
            result = a % b
            return f"{a} mod {b} = {result}"
        except Exception as e:
            return f"Error: {str(e)}"
    
    # Create calculator tools
    calc_tools = [
        Tool(name="multiply", description="Multiply two numbers. Use format: multiply(a, b)", func=lambda input_str: multiply_func(*map(float, input_str.split(',')))),
        Tool(name="add", description="Add two numbers. Use format: add(a, b)", func=lambda input_str: add_func(*map(float, input_str.split(',')))),
        Tool(name="subtract", description="Subtract two numbers. Use format: subtract(a, b)", func=lambda input_str: subtract_func(*map(float, input_str.split(',')))),
        Tool(name="divide", description="Divide two numbers. Use format: divide(a, b)", func=lambda input_str: divide_func(*map(float, input_str.split(',')))),
        Tool(name="modulus", description="Get modulus of two integers. Use format: modulus(a, b)", func=lambda input_str: modulus_func(*map(int, input_str.split(',')))),
    ]
    
    tools.extend(calc_tools)
    print(f"✅ Added {len(calc_tools)} calculator tools")
    
    # Hub stats tool
    try:
        from tools import get_hub_stats
        
        def hub_stats_func(author: str) -> str:
            """Get Hugging Face Hub statistics for an author"""
            try:
                return get_hub_stats(author)
            except Exception as e:
                return f"Hub stats error: {str(e)}"
        
        hub_tool = Tool(
            name="hub_stats",
            description="Get statistics for Hugging Face Hub models by author",
            func=hub_stats_func
        )
        tools.append(hub_tool)
        print("✅ Added Hub stats tool")
    except Exception as e:
        print(f"⚠️  Could not load Hub stats tool: {e}")
    
    # Python execution tool
    python_tool = create_python_execution_tool()
    tools.append(python_tool)
    print("✅ Added Python execution tool")
    
    print(f"🔧 Code Agent loaded {len(tools)} tools")
    return tools


def create_python_execution_tool() -> Tool:
    """Create a tool for executing Python code safely"""
    
    def execute_python_code(code: str) -> str:
        """
        Execute Python code in a controlled environment.
        
        Args:
            code: Python code to execute
            
        Returns:
            String containing the output or error message
        """
        # Create a string buffer to capture output
        output_buffer = io.StringIO()
        error_buffer = io.StringIO()
        
        # Prepare a safe execution environment
        safe_globals = {
            '__builtins__': {
                'print': lambda *args, **kwargs: print(*args, file=output_buffer, **kwargs),
                'len': len,
                'str': str,
                'int': int,
                'float': float,
                'list': list,
                'dict': dict,
                'set': set,
                'tuple': tuple,
                'range': range,
                'sum': sum,
                'max': max,
                'min': min,
                'abs': abs,
                'round': round,
                'sorted': sorted,
                'enumerate': enumerate,
                'zip': zip,
                'map': map,
                'filter': filter,
            }
        }
        
        # Allow common safe modules
        try:
            import math
            import statistics
            import datetime
            import json
            import re
            
            safe_globals.update({
                'math': math,
                'statistics': statistics,
                'datetime': datetime,
                'json': json,
                're': re,
            })
        except ImportError:
            pass
        
        try:
            # Execute the code
            with contextlib.redirect_stdout(output_buffer), \
                 contextlib.redirect_stderr(error_buffer):
                exec(code, safe_globals)
            
            # Get the output
            output = output_buffer.getvalue()
            error = error_buffer.getvalue()
            
            if error:
                return f"Error: {error}"
            elif output:
                return output.strip()
            else:
                return "Code executed successfully (no output)"
                
        except Exception as e:
            return f"Execution error: {str(e)}"
        finally:
            output_buffer.close()
            error_buffer.close()
    
    return Tool(
        name="python_execution",
        description="Execute Python code for calculations and data processing. Use for complex computations, data analysis, or when calculator tools are insufficient.",
        func=execute_python_code
    )


def load_code_prompt() -> str:
    """Load the code execution prompt"""
    try:
        with open("archive/prompts/execution_prompt.txt", "r") as f:
            return f.read()
    except FileNotFoundError:
        return """
You are a computational specialist focused on accurate calculations and code execution.

Your goals:
1. Perform mathematical calculations accurately
2. Write and execute Python code for complex computations
3. Process data and perform analysis as needed
4. Provide clear, numerical results

When handling computational tasks:
- Use calculator tools for basic arithmetic operations
- Use Python execution for complex calculations, data processing, or multi-step computations
- Show your work and intermediate steps
- Verify results when possible
- Handle edge cases and potential errors

Available tools:
- Calculator tools: add, subtract, multiply, divide, modulus
- Python execution: for complex computations and data analysis
- Hub stats tool: for Hugging Face model information

Format your response as:
### Computational Analysis
[Description of the approach]

### Calculations
[Step-by-step calculations or code]

### Results
[Final numerical results or outputs]
"""


def code_agent(state: Dict[str, Any]) -> Command:
    """
    Code Agent node that handles computational tasks and code execution.
    
    Returns Command with computational results appended to code_outputs.
    """
    
    print("🧮 Code Agent: Processing computational tasks...")
    
    try:
        # Get code execution prompt
        code_prompt = load_code_prompt()
        
        # Initialize LLM with tools
        llm = ChatGroq(
            model="llama-3.3-70b-versatile",
            temperature=0.1,  # Low temperature for accuracy in calculations
            max_tokens=2048
        )
        
        # Get computational tools
        tools = create_code_tools()
        
        # Bind tools to LLM
        llm_with_tools = llm.bind_tools(tools)
        
        # Create agent span for tracing
        with agent_span(
            "code",
            metadata={
                "tools_available": len(tools),
                "research_context_length": len(state.get("research_notes", "")),
                "user_id": state.get("user_id", "unknown"),
                "session_id": state.get("session_id", "unknown")
            }
        ) as span:
            
            # Extract user query and research context
            messages = state.get("messages", [])
            user_query = ""
            for msg in messages:
                if isinstance(msg, HumanMessage):
                    user_query = msg.content
                    break
            
            research_notes = state.get("research_notes", "")
            
            # Build computational request
            code_request = f"""
Please analyze the following question and perform any necessary calculations or code execution:

Question: {user_query}

Research Context:
{research_notes}

Current computational work: {len(state.get('code_outputs', ''))} characters already completed

Instructions:
1. Identify any computational or mathematical aspects of the question
2. Use appropriate tools for calculations or code execution
3. Show your work and intermediate steps
4. Provide clear, accurate results
5. If no computation is needed, state that clearly

Please perform all necessary calculations to help answer this question.
"""
            
            # Create messages for code execution
            code_messages = [
                SystemMessage(content=code_prompt),
                HumanMessage(content=code_request)
            ]
            
            # Get computational response
            response = llm_with_tools.invoke(code_messages)
            
            # Process the response - handle both tool calls and direct responses
            computation_results = []
            
            # Check if the LLM wants to use tools
            if hasattr(response, 'tool_calls') and response.tool_calls:
                print(f"🛠️  Executing {len(response.tool_calls)} computational operations")
                
                # Execute tool calls and collect results
                for tool_call in response.tool_calls:
                    try:
                        # Find the tool by name
                        tool = next((t for t in tools if t.name == tool_call['name']), None)
                        if tool:
                            # Handle different argument formats
                            args = tool_call.get('args', {})
                            if isinstance(args, dict):
                                # Convert dict args to string for simple tools
                                if len(args) == 1:
                                    arg_value = list(args.values())[0]
                                else:
                                    arg_value = ','.join(str(v) for v in args.values())
                            else:
                                arg_value = str(args)
                            
                            result = tool.func(arg_value)
                            computation_results.append(f"**{tool.name}**: {result}")
                        else:
                            computation_results.append(f"**{tool_call['name']}**: Tool not found")
                    except Exception as e:
                        print(f"⚠️  Tool {tool_call.get('name', 'unknown')} failed: {e}")
                        computation_results.append(f"**{tool_call.get('name', 'unknown')}**: Error - {str(e)}")
            
            # Compile computational results
            if computation_results:
                computational_findings = "\n\n".join(computation_results)
            else:
                # No tools used or tool calls failed, analyze if computation is needed
                computational_findings = response.content if hasattr(response, 'content') else str(response)
                
                # If the response looks like it should have used tools but didn't, try direct calculation
                if any(op in user_query.lower() for op in ['+', '-', '*', '/', 'calculate', 'compute', 'multiply', 'add', 'subtract', 'divide']):
                    print("🔧 Attempting direct calculation...")
                    
                    # Try to extract and solve simple mathematical expressions
                    import re
                    
                    # Look for simple math expressions
                    math_patterns = [
                        r'(\d+)\s*\+\s*(\d+)',  # addition
                        r'(\d+)\s*\*\s*(\d+)',  # multiplication  
                        r'(\d+)\s*-\s*(\d+)',   # subtraction
                        r'(\d+)\s*/\s*(\d+)',   # division
                    ]
                    
                    for pattern in math_patterns:
                        matches = re.findall(pattern, user_query)
                        if matches:
                            for match in matches:
                                a, b = int(match[0]), int(match[1])
                                if '+' in user_query:
                                    result = a + b
                                    computational_findings += f"\n\nDirect calculation: {a} + {b} = {result}"
                                elif '*' in user_query:
                                    result = a * b
                                    computational_findings += f"\n\nDirect calculation: {a} × {b} = {result}"
                                elif '-' in user_query:
                                    result = a - b
                                    computational_findings += f"\n\nDirect calculation: {a} - {b} = {result}"
                                elif '/' in user_query:
                                    result = a / b
                                    computational_findings += f"\n\nDirect calculation: {a} ÷ {b} = {result}"
                            break
            
            # Format computational results
            formatted_results = f"""
### Computational Analysis {state.get('loop_counter', 0) + 1}

{computational_findings}

---
"""
            
            print(f"🧮 Code Agent: Generated {len(formatted_results)} characters of computational results")
            
            # Update span with results if available
            if span:
                span.update_trace(metadata={
                    "computation_length": len(formatted_results),
                    "results_preview": formatted_results[:300] + "..."
                })
            
            # Return command to go back to lead agent
            return Command(
                goto="lead",
                update={
                    "code_outputs": state.get("code_outputs", "") + formatted_results
                }
            )
            
    except Exception as e:
        print(f"❌ Code Agent Error: {e}")
        
        # Return with error information
        error_result = f"""
### Computational Error
An error occurred during code execution: {str(e)}

"""
        return Command(
            goto="lead",
            update={
                "code_outputs": state.get("code_outputs", "") + error_result
            }
        )