""" 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.workflowAgentsRegistry 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.executorTimeout = int(APP_CONFIG.get("Agent_Coder_EXECUTION_TIMEOUT")) # seconds self.executionRetryLimit = int(APP_CONFIG.get("Agent_Coder_EXECUTION_RETRY")) # max retries self.tempDir = None def setDependencies(self, mydom=None): """Set external dependencies for the agent.""" self.mydom = mydom async def processTask(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, inputDocuments, outputSpecifications Returns: Dictionary with feedback and documents """ # 1. Extract task information prompt = task.get("prompt", "") inputDocuments = task.get("inputDocuments", []) outputSpecs = task.get("outputSpecifications", []) # Check if AI service is available if not self.mydom: 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 documentData = [] # For raw file data (for code execution) contentData = [] # For content data (later use) contentExtraction = [] # For AI-extracted data (for quick completion) for doc in inputDocuments: # 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 documentData if it exists docData = doc.get('data', '') if docData: isBase64 = True # Assume base64 encoded for document data documentData.append([filename, docData, isBase64]) # Process contents for different uses if doc.get('contents'): for content in doc.get('contents', []): contentName = content.get('name', 'unnamed') # For AI-extracted data (quick completion) if content.get('dataExtracted'): contentExtraction.append({ "filename": filename, "contentName": contentName, "contentData": content.get('dataExtracted', ''), "contentType": content.get('contentType', ''), "summary": content.get('summary', '') }) # For raw content data if content.get('data'): rawData = content.get('data', '') isBase64 = content.get('metadata', {}).get('base64Encoded', False) contentData.append({ "filename": filename, "contentName": contentName, "data": rawData, "isBase64": isBase64, "contentType": content.get('contentType', '') }) # Also add to documentData for code execution if not already added if not docData or docData != rawData: documentData.append([filename, rawData, isBase64]) # 3. Check if task can be completed without code execution quickCompletion = await self._checkQuickCompletion(prompt, contentExtraction, outputSpecs) if quickCompletion and quickCompletion.get("complete") == 1: logger.info("Task completed without code execution") return { "feedback": quickCompletion.get("prompt", "Task completed successfully."), "documents": quickCompletion.get("documents", []) } else: logger.debug(f"Code to generate, no quick check") # 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._generateCode(prompt) if not code: return { "feedback": "Failed to generate code for the task.", "documents": [] } # 5. Replace the placeholder with actual inputFiles data documentDataJson = repr(documentData) codeWithData = code.replace("inputFiles = \"=== JSONLOAD ===\"", f"inputFiles = {documentDataJson}") # 6. Execute code with retry logic retryCount = 0 maxRetries = self.executionRetryLimit executionHistory = [] while retryCount <= maxRetries: executionResult = self._executeCode(codeWithData, requirements) executionHistory.append({ "attempt": retryCount + 1, "code": codeWithData, "result": executionResult }) # Check if execution was successful if executionResult.get("success", False): logger.info(f"Code execution succeeded on attempt {retryCount + 1}") break # If we've reached max retries, exit the loop if retryCount >= maxRetries: logger.info(f"Reached maximum retry limit ({maxRetries}). Giving up.") break # Log the error and attempt to improve the code error = executionResult.get("error", "Unknown error") logger.info(f"Execution attempt {retryCount + 1} failed: {error}. Attempting to improve code.") # Generate improved code based on error improvedCode, improvedRequirements = await self._improveCode( originalCode=codeWithData, error=error, executionResult=executionResult, attempt=retryCount + 1 ) if improvedCode: codeWithData = improvedCode requirements = improvedRequirements logger.info(f"Code improved for retry {retryCount + 2}") else: logger.warning("Failed to improve code, using original code for retry") retryCount += 1 # 7. Process results and create output documents documents = [] # Always add the final code document documents.append({ "label": "generated_code.py", "content": codeWithData }) # Add execution history document executionHistoryStr = json.dumps(executionHistory, indent=2) documents.append({ "label": "execution_history.json", "content": executionHistoryStr }) # Create documents based on execution results if executionResult.get("success", False): resultData = executionResult.get("result") # Create documents based on output specifications if outputSpecs: for spec in outputSpecs: label = spec.get("label", "output.txt") # Extract content from result if available content = "" if isinstance(resultData, dict) and label in resultData: content = resultData[label] else: # Default to execution output content = executionResult.get("output", "") documents.append({ "label": label, "content": content }) else: # No output specs, create default output document documents.append({ "label": "execution_output.txt", "content": executionResult.get("output", "") }) if retryCount > 0: feedback = f"Code executed successfully after {retryCount + 1} attempts. Generated output files based on specifications." else: feedback = "Code executed successfully. Generated output files based on specifications." else: # Execution failed error = executionResult.get("error", "Unknown error") documents.append({ "label": "execution_error.txt", "content": f"Error executing code:\n\n{error}" }) if retryCount > 0: feedback = f"Error during code execution after {retryCount + 1} attempts: {error}" else: feedback = f"Error during code execution: {error}" return { "feedback": feedback, "documents": documents } async def _improveCode(self, originalCode: str, error: str, executionResult: Dict[str, Any], attempt: int) -> Tuple[str, List[str]]: """ Improve code based on execution error. Args: originalCode: The code that failed to execute error: The error message executionResult: Complete execution result dictionary attempt: Current attempt number Returns: Tuple of (improvedCode, requirements) """ # Create prompt for code improvement improvementPrompt = f""" Fix the following Python code that failed during execution. This is attempt {attempt} to fix the code. ORIGINAL CODE: {originalCode} ERROR MESSAGE: {error} STDOUT: {executionResult.get('output', '')} INSTRUCTIONS: 1. Fix all errors identified in the error message 2. Diagnose and fix any logical issues 3. Pay special attention to: - Type conversions and data handling - Error handling and edge cases - Resource management (file handles, etc.) - Syntax errors and typos 4. Keep the inputFiles handling logic intact 5. Maintain the same overall structure and purpose OUTPUT: - Your improved code MUST still define a 'result' variable as a dictionary - Each output file should be a key in the result dictionary - DO NOT remove the inputFiles assignment line structure REQUIREMENTS: Required packages should be specified as: # REQUIREMENTS: library==version,library2>=version - You may add/remove requirements as needed to fix the code Return ONLY Python code without explanations or markdown. """ # Call AI service messages = [ {"role": "system", "content": "You are an expert Python code debugger. Provide only fixed Python code without explanations or formatting."}, {"role": "user", "content": improvementPrompt} ] try: improvedContent = await self.mydom.callAi(messages, temperature=0.2) # Extract code and requirements improvedCode = self._cleanCode(improvedContent) # Extract requirements requirements = [] for line in improvedCode.split('\n'): if line.strip().startswith("# REQUIREMENTS:"): reqStr = line.replace("# REQUIREMENTS:", "").strip() requirements = [r.strip() for r in reqStr.split(',') if r.strip()] break return improvedCode, requirements except Exception as e: logger.error(f"Error improving code: {str(e)}") return None, [] async def _checkQuickCompletion(self, prompt: str, contentExtraction: List[Dict], outputSpecs: List[Dict]) -> Dict: """ Check if the task can be completed without writing and executing code. Args: prompt: The task prompt contentExtraction: List of extracted content data with contentName and dataExtracted outputSpecs: 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 contentExtraction or not outputSpecs: return None # Create a prompt for the AI to check if this can be completed directly specsJson = json.dumps(outputSpecs) dataJson = json.dumps(contentExtraction) checkPrompt = f""" Analyze this task and determine if it can be completed directly without writing code. TASK: {prompt} EXTRACTED DATA AVAILABLE: {dataJson} Each entry in the extracted data contains: - filename: The source file name - contentName: The specific content section name - contentData: The AI-extracted text from the content - contentType: The type of content (text, csv, etc.) - summary: A brief summary of the content REQUIRED OUTPUT: {specsJson} 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(f"Checking if task can be completed without code execution: {checkPrompt}") messages = [ {"role": "system", "content": "You are an AI assistant that determines if tasks require code execution. Reply with JSON only."}, {"role": "user", "content": checkPrompt} ] try: # Use a lower temperature for more deterministic response response = await self.mydom.callAi(messages, produceUserAnswer = True, temperature=0.1) # Parse response as JSON if response: try: # Find JSON in response if there's any text around it jsonStart = response.find('{') jsonEnd = response.rfind('}') + 1 if jsonStart >= 0 and jsonEnd > jsonStart: jsonStr = response[jsonStart:jsonEnd] result = json.loads(jsonStr) # 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 _generateCode(self, prompt: str) -> Tuple[str, List[str]]: """ Generate Python code from a prompt with the inputFiles placeholder. Args: prompt: The task prompt Returns: Tuple of (code, requirements) """ # Create prompt for code generation aiPrompt = f""" Generate Python code to solve the following task: TASK: {prompt} INPUT FILES: - 'inputFiles' variable is provided as [[filename, data, isBase64], ...] - For text files (isBase64=False): use data directly as string - For binary files (isBase64=True): use base64.b64decode(data) CODE QUALITY: - Use explicit type conversions where needed (int/float/str) - Implement feature detection, not version checks - Handle errors gracefully with appropriate fallbacks - Follow latest API conventions for libraries - Validate inputs before processing OUTPUT: - 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: inputFiles = "=== JSONLOAD ===" # DO NOT CHANGE THIS LINE REQUIREMENTS: Required packages should be specified as: # REQUIREMENTS: library==version,library2>=version - Specify exact versions for critical libraries - Use constraint operators (==,>=,<=) as needed Return ONLY Python code without explanations or markdown. """ # 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": aiPrompt} ] generatedContent = await self.mydom.callAi(messages, temperature=0.1) # Extract code and requirements code = self._cleanCode(generatedContent) # Extract requirements requirements = [] for line in code.split('\n'): if line.strip().startswith("# REQUIREMENTS:"): reqStr = line.replace("# REQUIREMENTS:", "").strip() requirements = [r.strip() for r in reqStr.split(',') if r.strip()] break return code, requirements def _executeCode(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.tempDir = tempfile.mkdtemp(prefix="code_exec_") venvPath = os.path.join(self.tempDir, "venv") # Create venv logger.debug(f"Creating virtual environment at {venvPath}") subprocess.run([sys.executable, "-m", "venv", venvPath], check=True, capture_output=True) # Get Python executable path pythonExe = os.path.join(venvPath, "Scripts", "python.exe") if os.name == 'nt' else os.path.join(venvPath, "bin", "python") # 2. Install requirements if provided if requirements: logger.info(f"Installing requirements: {requirements}") # Create requirements.txt reqFile = os.path.join(self.tempDir, "requirements.txt") with open(reqFile, "w") as f: f.write("\n".join(requirements)) x="\n".join(requirements) logger.info(f"Requirements file: {x}.") # Install requirements try: pipResult = subprocess.run( [pythonExe, "-m", "pip", "install", "-r", reqFile], capture_output=True, text=True, timeout=int(APP_CONFIG.get("Agent_Coder_INSTALL_TIMEOUT")) ) if pipResult.returncode != 0: logger.debug(f"Error installing requirements: {pipResult.stderr}") else: logger.debug(f"Requirements installed successfully") # Log installed packages if in debug mode if logger.isEnabledFor(logging.DEBUG): pipList = subprocess.run( [pythonExe, "-m", "pip", "list"], capture_output=True, text=True ) logger.debug(f"Installed packages:\n{pipList.stdout}") except Exception as e: logger.debug(f"Exception during requirements installation: {str(e)}") # 3. Write code to file codeFile = os.path.join(self.tempDir, "code.py") with open(codeFile, "w", encoding="utf-8") as f: f.write(code) # 4. Execute code logger.debug(f"Executing code with timeout of {self.executorTimeout} seconds. Code: {code}") process = subprocess.run( [pythonExe, codeFile], timeout=self.executorTimeout, capture_output=True, text=True ) # 5. Process results stdout = process.stdout stderr = process.stderr # Try to extract result from stdout resultData = 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: resultData = json.loads(line) logger.debug(f"Extracted result data from stdout: {type(resultData)}") 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": resultData, "exitCode": process.returncode } except subprocess.TimeoutExpired: logger.error(f"Execution timed out after {self.executorTimeout} seconds") return { "success": False, "output": "", "error": f"Execution timed out after {self.executorTimeout} seconds", "result": None, "exitCode": -1 } except Exception as e: logger.error(f"Execution error: {str(e)}") return { "success": False, "output": "", "error": f"Execution error: {str(e)}", "result": None, "exitCode": -1 } finally: # Clean up resources self._cleanupExecution() def _cleanupExecution(self): """Clean up temporary resources from code execution.""" if self.tempDir and os.path.exists(self.tempDir): try: logger.debug(f"Cleaning up temporary directory: {self.tempDir}") shutil.rmtree(self.tempDir) self.tempDir = None except Exception as e: logger.warning(f"Error cleaning up temp directory: {str(e)}") def _cleanCode(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') startIndex = 0 endIndex = len(lines) # Find start of actual code for i, line in enumerate(lines): if line.strip().startswith("inputFiles =") or line.strip().startswith("# REQUIREMENTS:"): startIndex = i break # Clean code cleanedCode = '\n'.join(lines[startIndex:endIndex]) return cleanedCode.strip() # Factory function for the Coder agent def getAgentCoder(): """Returns an instance of the Coder agent.""" return AgentCoder()