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 BaseTool, tool
from observability import agent_span, tool_span
from dotenv import load_dotenv

# Import LangChain-compatible tools
from langgraph_tools import get_code_tools

load_dotenv("env.local")


@tool("python_execution")
def python_execution_tool(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()


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
- Use Hugging Face Hub stats for model information
- 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
- Hugging Face Hub stats: for 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 using LangChain tools.
    
    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 (calculator tools + hub stats + python execution)
        code_tools = get_code_tools()
        code_tools.append(python_execution_tool)  # Add Python execution tool
        
        # Bind tools to LLM
        llm_with_tools = llm.bind_tools(code_tools)
        
        # Create agent span for tracing
        with agent_span(
            "code",
            metadata={
                "tools_available": len(code_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 tool calls if any
            computation_results = []
            if hasattr(response, 'tool_calls') and response.tool_calls:
                print(f"🛠️  Executing {len(response.tool_calls)} computational operations")
                
                for tool_call in response.tool_calls:
                    try:
                        # Find the tool by name
                        tool = next((t for t in code_tools if t.name == tool_call['name']), None)
                        if tool:
                            # Execute the tool
                            with tool_span(tool.name, metadata={"args": tool_call.get('args', {})}) as tool_span_ctx:
                                result = tool.invoke(tool_call.get('args', {}))
                                computation_results.append(f"**{tool.name}**: {result}")
                                if tool_span_ctx:
                                    tool_span_ctx.update_trace(output={"result": str(result)[:200] + "..."})
                        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)
                
                # Ask LLM to analyze the computational results
                analysis_request = f"""
Based on the computational results below, provide a structured analysis:

Original Question: {user_query}

Computational Results:
{computational_findings}

Please analyze these results and provide:
1. Summary of calculations performed
2. Key numerical findings
3. Interpretation of results
4. How these results help answer the original question

Structure your response clearly.
"""
                
                analysis_messages = [
                    SystemMessage(content=code_prompt),
                    HumanMessage(content=analysis_request)
                ]
                
                analysis_response = llm.invoke(analysis_messages)
                analysis_content = analysis_response.content if hasattr(analysis_response, 'content') else str(analysis_response)
                
                # Format final computational results
                formatted_results = f"""
### Computational Analysis {state.get('loop_counter', 0) + 1}

{analysis_content}

### Tool Results
{computational_findings}

---
"""
            else:
                # No tools were called, use the LLM response directly
                response_content = response.content if hasattr(response, 'content') else str(response)
                formatted_results = f"""
### Computational Analysis {state.get('loop_counter', 0) + 1}

{response_content}

---
"""
            
            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),
                    "tools_used": len(computation_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
            }
        )