""" AI Agent for executing and formatting code responses. This module provides functionality to execute Python code and format the results in a consistent way. It's designed to be used by the expert system for handling code-related tasks. """ import ast import sys import traceback from typing import Any, Dict, Optional, Tuple import json import logging # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) class CodeExecutionError(Exception): """Custom exception for code execution errors.""" pass class CodeAgent: """ Agent for executing and formatting Python code. This agent is responsible for: 1. Executing Python code in a controlled environment 2. Capturing and formatting the output 3. Handling errors gracefully 4. Providing rich output formatting """ def __init__(self): """Initialize the code agent with a clean environment.""" self.environment = self._create_clean_environment() def _create_clean_environment(self) -> Dict[str, Any]: """Create a clean execution environment with essential built-ins.""" # Start with a clean environment env = { '__builtins__': { # Safe built-ins 'abs': abs, 'all': all, 'any': any, 'ascii': ascii, 'bin': bin, 'bool': bool, 'bytearray': bytearray, 'bytes': bytes, 'chr': chr, 'complex': complex, 'dict': dict, 'divmod': divmod, 'enumerate': enumerate, 'filter': filter, 'float': float, 'format': format, 'frozenset': frozenset, 'hash': hash, 'hex': hex, 'int': int, 'iter': iter, 'len': len, 'list': list, 'map': map, 'max': max, 'min': min, 'next': next, 'oct': oct, 'ord': ord, 'pow': pow, 'range': range, 'repr': repr, 'reversed': reversed, 'round': round, 'set': set, 'slice': slice, 'sorted': sorted, 'str': str, 'sum': sum, 'tuple': tuple, 'zip': zip, # Constants 'True': True, 'False': False, 'None': None }, # Add common modules that are generally safe 'math': __import__('math'), 'random': __import__('random'), 'datetime': __import__('datetime'), 'json': __import__('json'), 'collections': __import__('collections'), 'itertools': __import__('itertools'), 'functools': __import__('functools'), 'operator': __import__('operator'), 're': __import__('re'), } return env def execute_code(self, code: str) -> Tuple[Any, Optional[str], bool]: """ Execute Python code and return the result, output, and success status. Args: code: The Python code to execute Returns: Tuple of (result, output, success) where: - result: The result of the last expression (or None) - output: Any captured stdout/stderr output - success: Boolean indicating if execution was successful """ # Redirect stdout and stderr from io import StringIO # Create a custom class to capture print output class PrintCapture(StringIO): def __init__(self): super().__init__() self.output = [] def write(self, text): self.output.append(text) return super().write(text) def getvalue(self): return ''.join(self.output) # Set up the capture output_capture = PrintCapture() old_stdout = sys.stdout old_stderr = sys.stderr sys.stdout = output_capture sys.stderr = output_capture try: # Parse the code to an AST tree = ast.parse(code) # If the last node is an expression, we'll capture its value result = None if tree.body and isinstance(tree.body[-1], ast.Expr): # Extract the last expression last_expr = tree.body.pop() # Compile and execute the remaining code if tree.body: exec( compile( ast.Module(body=tree.body, type_ignores=[]), '', 'exec' ), self.environment ) # Evaluate the last expression result = eval( compile( ast.Expression(body=last_expr.value), '', 'eval' ), self.environment ) else: # No expression to evaluate, just execute the code exec(code, self.environment) # Get the captured output output = output_capture.getvalue().strip() return result, output, True except Exception as e: # Capture the full traceback exc_type, exc_value, exc_traceback = sys.exc_info() tb_lines = traceback.format_exception(exc_type, exc_value, exc_traceback) error_msg = ''.join(tb_lines) # Get any output before the error output = output_buffer.getvalue() + '\n' + error_msg return None, output, False finally: # Restore stdout and stderr sys.stdout = old_stdout sys.stderr = old_stderr def format_code_response(self, code: str, result: Any, output: str, success: bool) -> str: """ Format the code execution results into a human-readable string. Args: code: The original code that was executed result: The result of the last expression (if any) output: Any captured stdout/stderr output success: Whether the execution was successful Returns: Formatted string with the code and its results """ response = [] # Add the code block response.append("```python") response.append(code) response.append("```") # Add the output if any if output.strip(): response.append("\n**Output:**") response.append("```") response.append(output.strip()) response.append("```") # Add the result if any if result is not None: response.append("\n**Result:**") response.append("```") response.append(str(result)) response.append("```") # Add status status = "✅ Success" if success else "❌ Error" response.append(f"\n**Status:** {status}") return '\n'.join(response) def process_code(self, code: str) -> str: """ Process and execute Python code, returning a formatted response. Args: code: The Python code to execute Returns: Formatted string with the code and its results """ try: # Clean and validate the code code = code.strip() if not code: return "Error: No code provided" # Execute the code result, output, success = self.execute_code(code) # Format the response return self.format_code_response(code, result, output, success) except Exception as e: error_msg = f"Error processing code: {str(e)}" logger.error(error_msg, exc_info=True) return f"Error: {error_msg}" def create_code_agent() -> 'CodeAgent': """ Factory function to create a new CodeAgent instance. Returns: A new instance of CodeAgent """ return CodeAgent() # Create a default instance for convenience agent = create_code_agent() def execute_code(code: str) -> str: """ Execute Python code and return the formatted results. This is a convenience function that uses the default agent instance. Args: code: The Python code to execute Returns: Formatted string with the code and its results """ return agent.process_code(code) if __name__ == "__main__": # Example usage test_code = """ # Calculate Fibonacci sequence def fibonacci(n): a, b = 0, 1 result = [] while a < n: result.append(a) a, b = b, a + b return result # Print first 10 Fibonacci numbers print("First 10 Fibonacci numbers:") fib_sequence = fibonacci(100) fib_sequence """ print(execute_code(test_code))