506 lines
No EOL
20 KiB
Python
506 lines
No EOL
20 KiB
Python
"""
|
|
Simple Coder Agent for execution of Python code.
|
|
"""
|
|
|
|
import logging
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import shutil
|
|
import sys
|
|
from typing import Dict, Any, List, Tuple
|
|
|
|
from modules.chat_registry import AgentBase
|
|
from modules.configuration import APP_CONFIG
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class AgentCoder(AgentBase):
|
|
"""Simplified Agent for developing and executing Python code with integrated executor"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the coder agent"""
|
|
super().__init__()
|
|
self.name = "coder"
|
|
self.description = "Develops and executes Python code for data processing and automation"
|
|
self.capabilities = [
|
|
"code_development",
|
|
"data_processing",
|
|
"file_processing",
|
|
"automation",
|
|
"code_execution"
|
|
]
|
|
|
|
# Executor settings
|
|
self.executor_timeout = APP_CONFIG.get("Agent_Coder_EXECUTION_TIMEOUT") # seconds
|
|
self.temp_dir = None
|
|
|
|
def set_dependencies(self, ai_service=None):
|
|
"""Set external dependencies for the agent."""
|
|
self.ai_service = ai_service
|
|
|
|
async def process_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Process a task and perform code development/execution.
|
|
First checks if the task can be completed without code execution,
|
|
then falls back to code generation if needed.
|
|
|
|
Args:
|
|
task: Task dictionary with prompt, input_documents, output_specifications
|
|
|
|
Returns:
|
|
Dictionary with feedback and documents
|
|
"""
|
|
# 1. Extract task information
|
|
prompt = task.get("prompt", "")
|
|
input_documents = task.get("input_documents", [])
|
|
output_specs = task.get("output_specifications", [])
|
|
|
|
# Check if AI service is available
|
|
if not self.ai_service:
|
|
logger.error("No AI service configured for the Coder agent")
|
|
return {
|
|
"feedback": "The Coder agent is not properly configured.",
|
|
"documents": []
|
|
}
|
|
|
|
# 2. Extract data from documents in separate categories
|
|
document_data = [] # For raw file data (for code execution)
|
|
content_data = [] # For content data (later use)
|
|
content_extraction = [] # For AI-extracted data (for quick completion)
|
|
|
|
for doc in input_documents:
|
|
# Create proper filename from name and ext
|
|
filename = f"{doc.get('name')}.{doc.get('ext')}" if doc.get('ext') else doc.get('name')
|
|
|
|
# Add main document data to document_data if it exists
|
|
doc_data = doc.get('data', '')
|
|
if doc_data:
|
|
is_base64 = True # Assume base64 encoded for document data
|
|
document_data.append([filename, doc_data, is_base64])
|
|
|
|
# Process contents for different uses
|
|
if doc.get('contents'):
|
|
for content in doc.get('contents', []):
|
|
content_name = content.get('name', 'unnamed')
|
|
|
|
# For AI-extracted data (quick completion)
|
|
if content.get('data_extracted'):
|
|
content_extraction.append({
|
|
"filename": filename,
|
|
"content_name": content_name,
|
|
"content_data": content.get('data_extracted', ''),
|
|
"content_type": content.get('content_type', ''),
|
|
"summary": content.get('summary', '')
|
|
})
|
|
|
|
# For raw content data
|
|
if content.get('data'):
|
|
raw_data = content.get('data', '')
|
|
is_base64 = content.get('metadata', {}).get('base64_encoded', False)
|
|
content_data.append({
|
|
"filename": filename,
|
|
"content_name": content_name,
|
|
"data": raw_data,
|
|
"is_base64": is_base64,
|
|
"content_type": content.get('content_type', '')
|
|
})
|
|
|
|
# Also add to document_data for code execution if not already added
|
|
if not doc_data or doc_data != raw_data:
|
|
document_data.append([filename, raw_data, is_base64])
|
|
|
|
# 3. Check if task can be completed without code execution
|
|
quick_completion = await self._check_quick_completion(prompt, content_extraction, output_specs)
|
|
|
|
if quick_completion and quick_completion.get("complete") == 1:
|
|
logger.info("Task completed without code execution")
|
|
return {
|
|
"feedback": quick_completion.get("prompt", "Task completed successfully."),
|
|
"documents": quick_completion.get("documents", [])
|
|
}
|
|
else:
|
|
logger.debug(f"Code to generate, quick check responded: {quick_completion.get("prompt", "(no answer)")}")
|
|
|
|
# If quick completion not possible, continue with code generation and execution
|
|
logger.info("Generating code to solve the task")
|
|
|
|
# 4. Generate code using AI
|
|
code, requirements = await self._generate_code(prompt, document_data)
|
|
|
|
if not code:
|
|
return {
|
|
"feedback": "Failed to generate code for the task.",
|
|
"documents": []
|
|
}
|
|
|
|
# 5. Replace the placeholder with actual input_files data
|
|
document_data_json = json.dumps(document_data)
|
|
code_with_data = code.replace("input_files = \"=== JSONLOAD ===\"", f"input_files = {document_data_json }")
|
|
|
|
# 6. Execute code and get results
|
|
execution_result = self._execute_code(code_with_data, requirements)
|
|
|
|
# 7. Process results and create output documents
|
|
documents = []
|
|
|
|
# Always add the code document
|
|
documents.append({
|
|
"label": "generated_code.py",
|
|
"content": code_with_data
|
|
})
|
|
|
|
# Create documents based on execution results
|
|
if execution_result.get("success", False):
|
|
result_data = execution_result.get("result")
|
|
|
|
# Create documents based on output specifications
|
|
if output_specs:
|
|
for spec in output_specs:
|
|
label = spec.get("label", "output.txt")
|
|
|
|
# Extract content from result if available
|
|
content = ""
|
|
if isinstance(result_data, dict) and label in result_data:
|
|
content = result_data[label]
|
|
else:
|
|
# Default to execution output
|
|
content = execution_result.get("output", "")
|
|
|
|
documents.append({
|
|
"label": label,
|
|
"content": content
|
|
})
|
|
else:
|
|
# No output specs, create default output document
|
|
documents.append({
|
|
"label": "execution_output.txt",
|
|
"content": execution_result.get("output", "")
|
|
})
|
|
|
|
feedback = "Code executed successfully. Generated output files based on specifications."
|
|
else:
|
|
# Execution failed
|
|
error = execution_result.get("error", "Unknown error")
|
|
documents.append({
|
|
"label": "execution_error.txt",
|
|
"content": f"Error executing code:\n\n{error}"
|
|
})
|
|
feedback = f"Error during code execution: {error}"
|
|
|
|
return {
|
|
"feedback": feedback,
|
|
"documents": documents
|
|
}
|
|
|
|
async def _check_quick_completion(self, prompt: str, content_extraction: List[Dict], output_specs: List[Dict]) -> Dict:
|
|
"""
|
|
Check if the task can be completed without writing and executing code.
|
|
|
|
Args:
|
|
prompt: The task prompt
|
|
content_extraction: List of extracted content data with content_name and data_extracted
|
|
output_specs: List of output specifications
|
|
|
|
Returns:
|
|
Dictionary with completion status and results, or None if no quick completion
|
|
"""
|
|
# If no data or no output specs, can't do a quick completion
|
|
if not content_extraction or not output_specs:
|
|
return None
|
|
|
|
# Create a prompt for the AI to check if this can be completed directly
|
|
specs_json = json.dumps(output_specs)
|
|
data_json = json.dumps(content_extraction)
|
|
|
|
check_prompt = f"""
|
|
Analyze this task and determine if it can be completed directly without writing code.
|
|
|
|
TASK:
|
|
{prompt}
|
|
|
|
EXTRACTED DATA AVAILABLE:
|
|
{data_json}
|
|
|
|
Each entry in the extracted data contains:
|
|
- filename: The source file name
|
|
- content_name: The specific content section name
|
|
- content_data: The AI-extracted text from the content
|
|
- content_type: The type of content (text, csv, etc.)
|
|
- summary: A brief summary of the content
|
|
|
|
REQUIRED OUTPUT:
|
|
{specs_json}
|
|
|
|
If the task can be completed directly with the available extracted data, respond with:
|
|
{{"complete": 1, "prompt": "Brief explanation of the solution", "documents": [
|
|
{{"label": "filename.ext", "content": "content here"}}
|
|
]}}
|
|
|
|
If code would be needed to properly complete this task, respond with:
|
|
{{"complete": 0, "prompt": "Explanation why code is needed"}}
|
|
|
|
Only return valid JSON. Your entire response must be parseable as JSON.
|
|
"""
|
|
|
|
# Call AI service
|
|
logger.debug("Checking if task can be completed without code execution")
|
|
messages = [
|
|
{"role": "system", "content": "You are an AI assistant that determines if tasks require code execution. Reply with JSON only."},
|
|
{"role": "user", "content": check_prompt}
|
|
]
|
|
|
|
try:
|
|
# Use a lower temperature for more deterministic response
|
|
response = await self.ai_service.call_api(messages, temperature=0.1)
|
|
|
|
# Parse response as JSON
|
|
if response:
|
|
try:
|
|
# Find JSON in response if there's any text around it
|
|
json_start = response.find('{')
|
|
json_end = response.rfind('}') + 1
|
|
|
|
if json_start >= 0 and json_end > json_start:
|
|
json_str = response[json_start:json_end]
|
|
result = json.loads(json_str)
|
|
|
|
# Check if this is a proper response
|
|
if "complete" in result:
|
|
return result
|
|
|
|
except json.JSONDecodeError:
|
|
logger.debug("Failed to parse quick completion response as JSON")
|
|
pass
|
|
except Exception as e:
|
|
logger.debug(f"Error during quick completion check: {str(e)}")
|
|
|
|
# Default to requiring code execution
|
|
return None
|
|
|
|
async def _generate_code(self, prompt: str, input_files: List) -> Tuple[str, List[str]]:
|
|
"""
|
|
Generate Python code from a prompt with the input_files placeholder.
|
|
|
|
Args:
|
|
prompt: The task prompt
|
|
input_files: List of [filename, data, is_base64] items
|
|
|
|
Returns:
|
|
Tuple of (code, requirements)
|
|
"""
|
|
# Create prompt for code generation
|
|
ai_prompt = f"""
|
|
Generate Python code to solve the following task:
|
|
|
|
TASK:
|
|
{prompt}
|
|
|
|
IMPORTANT:
|
|
- All input files are provided in the 'input_files' variable as a list of [filename, data, is_base64].
|
|
- The 'input_files' variable is already defined at the top of your code, DO NOT modify it.
|
|
- For each file, you can access:
|
|
- filename: The name of the file (e.g., "image.png")
|
|
- data: The content of the file (base64 encoded or plain text)
|
|
- is_base64: Boolean flag indicating if the data is base64 encoded
|
|
|
|
- To use a file's data:
|
|
- For text files (when is_base64=False): Use the data directly as a string
|
|
- For binary files (when is_base64=True): Use base64.b64decode(data) to get bytes
|
|
|
|
- Do not perform any additional base64 detection - rely on the is_base64 flag
|
|
|
|
- Your code MUST define a 'result' variable as a dictionary to store outputs.
|
|
- Each output file should be a key in the result dictionary.
|
|
- For example: result = {{"output.txt": "output text", "results.json": json_string}}
|
|
|
|
Your code must start with:
|
|
input_files = "=== JSONLOAD ===" # DO NOT CHANGE THIS LINE
|
|
|
|
Required packages should be specified as:
|
|
# REQUIREMENTS: package1,package2,package3
|
|
|
|
Return ONLY Python code without explanations or markdown formatting.
|
|
"""
|
|
|
|
# Call AI service
|
|
messages = [
|
|
{"role": "system", "content": "You are a Python code generator. Provide only valid Python code without explanations or formatting."},
|
|
{"role": "user", "content": ai_prompt}
|
|
]
|
|
|
|
generated_content = await self.ai_service.call_api(messages, temperature=0.1)
|
|
|
|
# Extract code and requirements
|
|
code = self._clean_code(generated_content)
|
|
|
|
# Extract requirements
|
|
requirements = []
|
|
for line in code.split('\n'):
|
|
if line.strip().startswith("# REQUIREMENTS:"):
|
|
req_str = line.replace("# REQUIREMENTS:", "").strip()
|
|
requirements = [r.strip() for r in req_str.split(',') if r.strip()]
|
|
break
|
|
|
|
return code, requirements
|
|
|
|
def _execute_code(self, code: str, requirements: List[str] = None) -> Dict[str, Any]:
|
|
"""
|
|
Execute Python code in a virtual environment.
|
|
Integrated executor functionality.
|
|
|
|
Args:
|
|
code: Python code to execute
|
|
requirements: List of required packages
|
|
|
|
Returns:
|
|
Execution result dictionary
|
|
"""
|
|
try:
|
|
# 1. Create temp directory and virtual environment
|
|
self.temp_dir = tempfile.mkdtemp(prefix="code_exec_")
|
|
venv_path = os.path.join(self.temp_dir, "venv")
|
|
|
|
# Create venv
|
|
logger.debug(f"Creating virtual environment at {venv_path}")
|
|
subprocess.run([sys.executable, "-m", "venv", venv_path],
|
|
check=True, capture_output=True)
|
|
|
|
# Get Python executable path
|
|
python_exe = os.path.join(venv_path, "Scripts", "python.exe") if os.name == 'nt' else os.path.join(venv_path, "bin", "python")
|
|
|
|
# 2. Install requirements if provided
|
|
if requirements:
|
|
logger.debug(f"Installing requirements: {requirements}")
|
|
|
|
# Create requirements.txt
|
|
req_file = os.path.join(self.temp_dir, "requirements.txt")
|
|
with open(req_file, "w") as f:
|
|
f.write("\n".join(requirements))
|
|
|
|
# Install requirements
|
|
try:
|
|
pip_result = subprocess.run(
|
|
[python_exe, "-m", "pip", "install", "-r", req_file],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=APP_CONFIG.get("Agent_Coder_INSTALL_TIMEOUT")
|
|
)
|
|
if pip_result.returncode != 0:
|
|
logger.debug(f"Error installing requirements: {pip_result.stderr}")
|
|
else:
|
|
logger.debug(f"Requirements installed successfully")
|
|
# Log installed packages if in debug mode
|
|
if logger.isEnabledFor(logging.DEBUG):
|
|
pip_list = subprocess.run(
|
|
[python_exe, "-m", "pip", "list"],
|
|
capture_output=True,
|
|
text=True
|
|
)
|
|
logger.debug(f"Installed packages:\n{pip_list.stdout}")
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Exception during requirements installation: {str(e)}")
|
|
|
|
# 3. Write code to file
|
|
code_file = os.path.join(self.temp_dir, "code.py")
|
|
with open(code_file, "w", encoding="utf-8") as f:
|
|
f.write(code)
|
|
|
|
# 4. Execute code
|
|
logger.debug(f"Executing code with timeout of {self.executor_timeout} seconds. Code: {code}")
|
|
process = subprocess.run(
|
|
[python_exe, code_file],
|
|
timeout=self.executor_timeout,
|
|
capture_output=True,
|
|
text=True
|
|
)
|
|
|
|
# 5. Process results
|
|
stdout = process.stdout
|
|
stderr = process.stderr
|
|
|
|
# Try to extract result from stdout
|
|
result_data = None
|
|
if process.returncode == 0:
|
|
try:
|
|
# Find the last line that might be JSON
|
|
for line in reversed(stdout.strip().split('\n')):
|
|
line = line.strip()
|
|
if line and line[0] in '{[' and line[-1] in '}]':
|
|
try:
|
|
result_data = json.loads(line)
|
|
logger.debug(f"Extracted result data from stdout: {type(result_data)}")
|
|
break
|
|
except json.JSONDecodeError:
|
|
continue
|
|
except Exception as e:
|
|
logger.debug(f"Error extracting result from stdout: {str(e)}")
|
|
|
|
# Create result dictionary
|
|
return {
|
|
"success": process.returncode == 0,
|
|
"output": stdout,
|
|
"error": stderr if process.returncode != 0 else "",
|
|
"result": result_data,
|
|
"exit_code": process.returncode
|
|
}
|
|
|
|
except subprocess.TimeoutExpired:
|
|
logger.error(f"Execution timed out after {self.executor_timeout} seconds")
|
|
return {
|
|
"success": False,
|
|
"output": "",
|
|
"error": f"Execution timed out after {self.executor_timeout} seconds",
|
|
"result": None,
|
|
"exit_code": -1
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Execution error: {str(e)}")
|
|
return {
|
|
"success": False,
|
|
"output": "",
|
|
"error": f"Execution error: {str(e)}",
|
|
"result": None,
|
|
"exit_code": -1
|
|
}
|
|
finally:
|
|
# Clean up resources
|
|
self._cleanup_execution()
|
|
|
|
def _cleanup_execution(self):
|
|
"""Clean up temporary resources from code execution."""
|
|
if self.temp_dir and os.path.exists(self.temp_dir):
|
|
try:
|
|
logger.debug(f"Cleaning up temporary directory: {self.temp_dir}")
|
|
shutil.rmtree(self.temp_dir)
|
|
self.temp_dir = None
|
|
except Exception as e:
|
|
logger.warning(f"Error cleaning up temp directory: {str(e)}")
|
|
|
|
def _clean_code(self, code: str) -> str:
|
|
"""Remove any markdown formatting or explanations."""
|
|
# Remove code block markers
|
|
code = code.replace("```python", "").replace("```", "")
|
|
|
|
# Remove explanations before or after code
|
|
lines = code.strip().split('\n')
|
|
start_index = 0
|
|
end_index = len(lines)
|
|
|
|
# Find start of actual code
|
|
for i, line in enumerate(lines):
|
|
if line.strip().startswith("input_files =") or line.strip().startswith("# REQUIREMENTS:"):
|
|
start_index = i
|
|
break
|
|
|
|
# Clean code
|
|
cleaned_code = '\n'.join(lines[start_index:end_index])
|
|
return cleaned_code.strip()
|
|
|
|
|
|
# Factory function for the Coder agent
|
|
def get_coder_agent():
|
|
"""Returns an instance of the Coder agent."""
|
|
return AgentCoder() |