File size: 6,377 Bytes
5ddeed7 5196d55 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
"""
Code execution tool for the forensic agent.
This allows the agent to execute Python code dynamically, similar to ChatGPT's code interpreter.
Useful for custom image analysis, zooming, cropping, statistical analysis, etc.
"""
import io
import sys
import os
import traceback
from typing import Dict, Any, Optional
from pathlib import Path
import base64
import json
try:
from PIL import Image
import numpy as np
except ImportError:
Image = None
np = None
def execute_python_code(code: str, image_path: Optional[str] = None, context: Optional[Dict[str, Any]] = None) -> str:
"""
Execute Python code in a sandboxed environment with image processing capabilities.
Args:
code: Python code to execute
image_path: Optional path to the current image being analyzed
context: Optional dictionary of variables to make available to the code
Returns:
String result of code execution (stdout + return value if any)
"""
# Capture stdout
old_stdout = sys.stdout
sys.stdout = captured_output = io.StringIO()
# Capture stderr
old_stderr = sys.stderr
sys.stderr = captured_error = io.StringIO()
try:
# Create a safe execution namespace
namespace = {
'__builtins__': {
'abs': abs, 'all': all, 'any': any, 'bool': bool, 'dict': dict,
'enumerate': enumerate, 'float': float, 'int': int, 'len': len,
'list': list, 'max': max, 'min': min, 'print': print, 'range': range,
'round': round, 'set': set, 'sorted': sorted, 'str': str, 'sum': sum,
'tuple': tuple, 'type': type, 'zip': zip,
# Math functions
'abs': abs, 'round': round, 'min': min, 'max': max, 'sum': sum,
},
'np': np,
'Image': Image,
'Path': Path,
'base64': base64,
'json': json,
}
# Add image path if provided
if image_path:
namespace['image_path'] = image_path
namespace['current_image_path'] = image_path
# Load image if PIL is available
if Image:
try:
img = Image.open(image_path)
namespace['image'] = img
namespace['img'] = img
# Also provide as numpy array if numpy is available
if np:
img_array = np.array(img)
namespace['image_array'] = img_array
namespace['img_array'] = img_array
except Exception as e:
namespace['image_load_error'] = str(e)
# Add any context variables
if context:
namespace.update(context)
# Execute the code
exec(code, namespace)
# Check if there's a return value
result_value = namespace.get('result', None)
# Get captured output
stdout_output = captured_output.getvalue()
stderr_output = captured_error.getvalue()
# Combine outputs
output_parts = []
if stdout_output:
output_parts.append(stdout_output)
if stderr_output:
output_parts.append(f"STDERR:\n{stderr_output}")
if result_value is not None:
output_parts.append(f"\nReturn value: {result_value}")
return '\n'.join(output_parts) if output_parts else "Code executed successfully (no output)"
except Exception as e:
error_traceback = traceback.format_exc()
return f"Error executing code:\n{error_traceback}"
finally:
# Restore stdout/stderr
sys.stdout = old_stdout
sys.stderr = old_stderr
def run_code_interpreter(input_str: str) -> str:
"""
LangChain tool wrapper for code execution.
Input format: JSON string with 'code' and optionally 'image_path'
Example: '{"code": "print(image.size)", "image_path": "path/to/image.jpg"}'
Or simple string: just the code (will try to extract image_path from context if available)
The agent should include the image_path in the JSON if available from the analysis context.
"""
try:
# Try to parse as JSON first
try:
params = json.loads(input_str)
code = params.get('code', '')
image_path = params.get('image_path')
context = params.get('context', {})
except (json.JSONDecodeError, AttributeError):
# If not JSON, treat as plain code string
code = input_str
image_path = None
context = {}
# Try to extract image_path from the code string if it mentions it
# This is a fallback - ideally the agent should pass it in JSON
import re
path_match = re.search(r'image_path\s*=\s*["\']([^"\']+)["\']', code)
if path_match:
image_path = path_match.group(1)
if not code.strip():
return "Error: No code provided. Provide Python code to execute.\n" \
"Format: {\"code\": \"your_python_code_here\", \"image_path\": \"path/to/image.jpg\"}\n" \
"Or: just the Python code as a string (image_path should be in agent context)."
# If no image_path provided, try to find it from common patterns in code
if not image_path:
import re
# Look for image_path variable assignment or usage
path_patterns = [
r'image_path\s*=\s*["\']([^"\']+)["\']',
r'["\']([^"\']+\.(jpg|jpeg|png|gif|bmp))["\']',
]
for pattern in path_patterns:
match = re.search(pattern, code, re.IGNORECASE)
if match:
potential_path = match.group(1)
if os.path.exists(potential_path):
image_path = potential_path
break
# Execute the code
result = execute_python_code(code, image_path=image_path, context=context)
return result
except Exception as e:
return f"Error in code interpreter: {str(e)}\n{traceback.format_exc()}"
|