""" Simple Coder Agent for execution of Python code. Modified to pass expected output document names to the generated 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. Enhanced to ensure all generated documents are included in output. 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, outputSpecs) 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, outputSpecs=outputSpecs ) 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(self.formatAgentDocumentOutput("generated_code.py", codeWithData, "text/plain")) # Add execution history document executionHistoryStr = json.dumps(executionHistory, indent=2) documents.append(self.formatAgentDocumentOutput("execution_history.json", executionHistoryStr, "application/json")) # Enhanced result handling: Create documents based on execution results - fixed for proper content extraction if executionResult.get("success", False): resultData = executionResult.get("result") # Process results from the result dictionary if available if isinstance(resultData, dict): # First, create a mapping of expected output labels to their specs expectedOutputs = {spec.get("label"): spec for spec in outputSpecs} createdOutputs = set() for label, result_item in resultData.items(): # Check if result follows the expected structure with nested content if isinstance(result_item, dict) and "content" in result_item: # Extract values from the properly structured result content = result_item.get("content", "") # Extract the inner content base64Encoded = result_item.get("base64Encoded", False) contentType = result_item.get("contentType", "text/plain") # Check if this label matches one of our expected output documents # If not, but we haven't created all expected outputs yet, try to map it finalLabel = label if label not in expectedOutputs and len(expectedOutputs) > 0: # Find an unused expected output label for expectedLabel in expectedOutputs: if expectedLabel not in createdOutputs: logger.warning(f"Remapping output '{label}' to expected '{expectedLabel}'") finalLabel = expectedLabel break # Create document by passing only the content to formatAgentDocumentOutput doc = self.formatAgentDocumentOutput(finalLabel, content, contentType) # Override the base64Encoded flag with the value from the result # This is needed since formatAgentDocumentOutput might determine a different value if isinstance(base64Encoded, bool): doc["base64Encoded"] = base64Encoded documents.append(doc) createdOutputs.add(finalLabel) logger.info(f"Created document from result: {finalLabel} ({contentType}, base64={base64Encoded})") else: # Not properly structured - log warning logger.warning(f"Skipping improperly formatted result for '{label}'. Results must include 'content' field.") else: # No result dictionary found logger.warning("No valid result dictionary found or it's not properly formatted") # If no valid documents were created from the result dictionary but we have output specifications if len(documents) <= 2 and outputSpecs: # Only code.py and history.json exist logger.warning("No valid documents created from result dictionary, using execution output for specifications") # Default to execution output output = executionResult.get("output", "") for spec in outputSpecs: label = spec.get("label", "output.txt") # Create basic document from output doc = self.formatAgentDocumentOutput(label, output, "text/plain") documents.append(doc) logger.info(f"Created document from output specification: {label}") if retryCount > 0: feedback = f"Code executed successfully after {retryCount + 1} attempts. Generated {len(documents) - 2} output files." else: feedback = f"Code executed successfully. Generated {len(documents) - 2} output files." else: # Execution failed error = executionResult.get("error", "Unknown error") documents.append(self.formatAgentDocumentOutput("execution_error.txt", f"Error executing code:\n\n{error}", "text/plain")) 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, outputSpecs: List[Dict[str, Any]] = None) -> Tuple[str, List[str]]: """ Improve code based on execution error. Enhanced to maintain proper output handling with correct document structure. Args: originalCode: The code that failed to execute error: The error message executionResult: Complete execution result dictionary attempt: Current attempt number outputSpecs: List of expected output specifications Returns: Tuple of (improvedCode, requirements) """ # Create a string with output specifications to be included in the prompt outputSpecsStr = "" if outputSpecs: outputSpecsStr = "\nEXPECTED OUTPUT DOCUMENTS:\n" for i, spec in enumerate(outputSpecs, 1): label = spec.get("label", f"output{i}.txt") description = spec.get("description", "") outputSpecsStr += f"{i}. {label} - {description}\n" # 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', '')} {outputSpecsStr} 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 REQUIREMENTS (VERY IMPORTANT): - Your code MUST define a 'result' variable as a dictionary to store ALL outputs - The key for each entry MUST be the full filename with extension (e.g., "output.txt") - The value for each entry MUST be a dictionary with the following structure: {{ "content": string, # The actual content (text or base64-encoded string) "base64Encoded": boolean, # Set to true for binary data, false for text data "contentType": string # MIME type of the content (e.g., "text/plain", "application/json") }} - Example result dictionary: result = {{ "output.txt": {{ "content": "This is text content", "base64Encoded": False, "contentType": "text/plain" }}, "chart.png": {{ "content": "base64encodedstring...", "base64Encoded": True, "contentType": "image/png" }} }} - NEVER write files to disk using open() or similar methods - use the result dictionary instead JSON OUTPUT (CRITICAL): - After creating the result dictionary, you MUST print it as JSON to stdout - Make sure your code includes: print(json.dumps(result)) as the final line - This printed JSON is how the system captures your result 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. Ensure all generated files are included in the 'result' dictionary and that result is printed as JSON with print(json.dumps(result))."}, {"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, outputSpecs: List[Dict[str, Any]] = None) -> Tuple[str, List[str]]: """ Generate Python code from a prompt with the inputFiles placeholder. Enhanced to emphasize proper result output handling with correct document structure. Args: prompt: The task prompt outputSpecs: List of expected output specifications Returns: Tuple of (code, requirements) """ # Create a string with output specifications to be included in the prompt outputSpecsStr = "" if outputSpecs: outputSpecsStr = "\nEXPECTED OUTPUT DOCUMENTS:\n" for i, spec in enumerate(outputSpecs, 1): label = spec.get("label", f"output{i}.txt") description = spec.get("description", "") outputSpecsStr += f"{i}. {label} - {description}\n" # Create improved prompt for code generation aiPrompt = f""" Generate Python code to solve the following task: TASK: {prompt} {outputSpecsStr} 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) OUTPUT REQUIREMENTS (VERY IMPORTANT): - Your code MUST define a 'result' variable as a dictionary to store ALL outputs - The key for each entry MUST be the full filename with extension (e.g., "output.txt") - The value for each entry MUST be a dictionary with the following structure: {{ "content": string, # The actual content (text or base64-encoded string) "base64Encoded": boolean, # Set to true for binary data, false for text data "contentType": string # MIME type of the content (e.g., "text/plain", "application/json") }} - Example result dictionary: result = {{ "output.txt": {{ "content": "This is text content", "base64Encoded": False, "contentType": "text/plain" }}, "chart.png": {{ "content": "base64encodedstring...", "base64Encoded": True, "contentType": "image/png" }} }} - NEVER write files to disk using open() or similar methods - use the result dictionary instead - If you generate any charts, reports, or visualizations, ensure they are properly encoded and included IMPORTANT - USE EXACT OUTPUT FILENAMES: - You MUST use the EXACT filenames specified in EXPECTED OUTPUT DOCUMENTS section - The key in the result dictionary must match these filenames precisely - If no output documents are specified, use appropriate descriptive filenames JSON OUTPUT (CRITICAL): - After creating the result dictionary, you MUST print it as JSON to stdout using json.dumps() - Add these lines at the end of your code: import json # if not already imported print(json.dumps(result)) - This printed JSON is how the system captures your result - Make sure this is the last thing your code prints BINARY DATA HANDLING: - For binary content (images, PDFs, etc.), convert to base64 string and set base64Encoded=True - For text content (text, JSON, HTML, etc.), use plain string and set base64Encoded=False - Use appropriate MIME types for different content types 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 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. Always output the result dictionary as JSON using print(json.dumps(result)) at the end of your code."}, {"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 with enhanced result extraction. 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 jsonLines = [] for line in stdout.strip().split('\n'): line = line.strip() if line and line[0] in '{[' and line[-1] in '}]': try: parsed = json.loads(line) jsonLines.append((line, parsed)) except json.JSONDecodeError: continue # Use the last valid JSON that appears to be a dictionary if jsonLines: for line, parsed in reversed(jsonLines): if isinstance(parsed, dict): resultData = parsed logger.debug(f"Extracted result data from stdout: {type(resultData)}") break except Exception as e: logger.debug(f"Error extracting result from stdout: {str(e)}") # Enhanced logging of what was found if resultData: logger.info(f"Found result dictionary with {len(resultData)} entries: {list(resultData.keys())}") else: logger.warning("No result dictionary found in output") # 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()