gateway/modules/agents/agentCoder.py

"""
Coder agent for generating and executing code.
Provides code generation, execution, and improvement capabilities.
"""

import logging
from typing import Dict, Any, List, Tuple
import json
import os
import sys
import subprocess
import tempfile
import shutil
import venv
import importlib.util
from datetime import datetime

from modules.workflow.agentBase import AgentBase
from modules.shared.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.label = "Developer and Code Executor"
        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, serviceBase=None):
        """Set external dependencies for the agent."""
        self.setService(serviceBase)

    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.service or not self.service.base:
            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": []
            }
        # Store the original code without document data
        original_clean_code = code  # Save clean code for later use in improvement

        # 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=original_clean_code,  # Use clean code without document data
                error=error,
                executionResult=executionResult,
                attempt=retryCount + 1,
                outputSpecs=outputSpecs
            )
                        
            if improvedCode:
                # Inject document data into improved code
                original_clean_code = improvedCode  # Update clean code for next potential improvement
                codeWithData = improvedCode.replace("inputFiles = \"=== JSONLOAD ===\"", f"inputFiles = {documentDataJson}")
                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. If there is a requirements error for missing or failes modules, then create alternate code with other modules
3. Diagnose and fix any logical issues
4. Pay special attention to:
- Type conversions and data handling
- Error handling and edge cases
- Resource management (file handles, etc.)
- Syntax errors and typos
5. Keep the inputFiles handling logic intact
6. 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.service.base.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.service.base.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.service.base.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 _executeCodeProd(self, code: str, requirements: List[str] = None) -> Dict[str, Any]:
        """
        Execute Python code in Azure environment using the antenv interpreter.
        Optimized for production use in Azure Web App environment where venv creation fails.
        
        Args:
            code: Python code to execute
            requirements: List of required packages
            
        Returns:
            Execution result dictionary
        """
        try:
            # 1. Create temp directory for code files
            self.tempDir = tempfile.mkdtemp(prefix="code_exec_")
            
            # Try different possible paths to find the antenv Python interpreter
            possible_python_paths = [
                "/home/site/wwwroot/antenv/bin/python",
                "/antenv/bin/python", 
                "/tmp/8dd8c226509f116/antenv/bin/python",  # Path from your error logs
                sys.executable  # Fallback to system Python
            ]
            
            pythonExe = None
            for path in possible_python_paths:
                if os.path.exists(path):
                    pythonExe = path
                    logger.info(f"Found Python interpreter at: {pythonExe}")
                    break
                    
            if not pythonExe:
                logger.error("Could not find a valid Python interpreter in Azure environment")
                return {
                    "success": False,
                    "output": "",
                    "error": "Could not find a valid Python interpreter in Azure environment",
                    "result": None,
                    "exitCode": -1
                }
            
            # 2. Install requirements to a temporary user directory if provided
            if requirements:
                logger.info(f"Installing requirements in Azure environment: {requirements}")
                
                # Create requirements.txt
                reqFile = os.path.join(self.tempDir, "requirements.txt")
                with open(reqFile, "w") as f:
                    f.write("\n".join(requirements))
                
                # Set up a custom PYTHONUSERBASE to isolate package installations
                custom_user_base = os.path.join(self.tempDir, "pip_packages")
                os.makedirs(custom_user_base, exist_ok=True)
                
                env = os.environ.copy()
                env["PYTHONUSERBASE"] = custom_user_base
                
                # Install requirements to the custom user directory
                try:
                    pipResult = subprocess.run(
                        [pythonExe, "-m", "pip", "install", "--user", "-r", reqFile],
                        capture_output=True,
                        text=True,
                        env=env,
                        timeout=int(APP_CONFIG.get("Agent_Coder_INSTALL_TIMEOUT"))
                    )
                    
                    if pipResult.returncode != 0:
                        logger.warning(f"Error installing requirements in Azure: {pipResult.stderr}")
                    else:
                        logger.info(f"Requirements installed successfully to {custom_user_base}")
                        
                    # Try to find the site-packages directory
                    import glob
                    site_packages = os.path.join(custom_user_base, "lib", "python*", "site-packages")
                    site_packages_paths = glob.glob(site_packages)
                    
                    if site_packages_paths:
                        env["PYTHONPATH"] = os.pathsep.join([site_packages_paths[0], env.get("PYTHONPATH", "")])
                        logger.info(f"Added {site_packages_paths[0]} to PYTHONPATH")
                    else:
                        # Alternative paths for different Python versions
                        alt_site_packages = os.path.join(custom_user_base, "site-packages")
                        if os.path.exists(alt_site_packages):
                            env["PYTHONPATH"] = os.pathsep.join([alt_site_packages, env.get("PYTHONPATH", "")])
                            logger.info(f"Added {alt_site_packages} to PYTHONPATH")
                except Exception as e:
                    logger.warning(f"Exception during requirements installation in Azure: {str(e)}")
            else:
                env = os.environ.copy()
                
            # 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 with the modified environment
            logger.debug(f"Executing code in Azure environment with timeout of {self.executorTimeout} seconds")
            process = subprocess.run(
                [pythonExe, codeFile],
                timeout=self.executorTimeout,
                capture_output=True,
                text=True,
                env=env
            )
            
            # 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 in Azure 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 in Azure environment: {str(e)}")
            return {
                "success": False,
                "output": "",
                "error": f"Execution error in Azure environment: {str(e)}",
                "result": None,
                "exitCode": -1
            }
        finally:
            # Clean up resources
            self._cleanupExecution()

    def _executeCodeVenv(self, code: str, requirements: List[str] = None) -> Dict[str, Any]:
        """
        Execute Python code in a virtual environment.
        Original implementation with venv creation for non-Azure environments.
        
        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}")
            
            try:
                # First try with sys.executable - the standard approach
                subprocess.run([sys.executable, "-m", "venv", venvPath], 
                            check=True, capture_output=True, timeout=60)
                logger.debug("Virtual environment created successfully with sys.executable")
            except (subprocess.SubprocessError, subprocess.CalledProcessError) as e:
                logger.warning(f"Failed to create venv with sys.executable: {str(e)}")
                
                # Fallback method 1: Try with explicit 'python3' command
                try:
                    logger.debug("Trying to create virtual environment with python3 command")
                    subprocess.run(["python3", "-m", "venv", venvPath], 
                                check=True, capture_output=True, timeout=60)
                    logger.debug("Virtual environment created successfully with python3")
                except (subprocess.SubprocessError, subprocess.CalledProcessError) as e:
                    logger.warning(f"Failed to create venv with python3: {str(e)}")
                    
                    # Fallback method 2: Try with virtualenv instead of venv
                    try:
                        logger.debug("Trying to create virtual environment with virtualenv module")
                        subprocess.run([sys.executable, "-m", "pip", "install", "virtualenv"], 
                                    check=False, capture_output=True, timeout=60)
                        subprocess.run([sys.executable, "-m", "virtualenv", venvPath], 
                                    check=True, capture_output=True, timeout=60)
                        logger.debug("Virtual environment created successfully with virtualenv")
                    except (subprocess.SubprocessError, subprocess.CalledProcessError) as e:
                        # If all methods fail, raise an exception
                        error_msg = f"Failed to create virtual environment with all methods: {str(e)}"
                        logger.error(error_msg)
                        raise RuntimeError(error_msg)
            
            # Get Python executable path - adjusted for OS
            if os.name == 'nt':  # Windows
                pythonExe = os.path.join(venvPath, "Scripts", "python.exe")
            else:  # Linux/Mac
                pythonExe = os.path.join(venvPath, "bin", "python")
            
            # Verify python executable exists
            if not os.path.exists(pythonExe):
                # Try to find it
                if os.name == 'nt':
                    possible_paths = [
                        os.path.join(venvPath, "Scripts", "python.exe"),
                        os.path.join(venvPath, "Scripts", "python")
                    ]
                else:
                    possible_paths = [
                        os.path.join(venvPath, "bin", "python"),
                        os.path.join(venvPath, "bin", "python3")
                    ]
                
                for path in possible_paths:
                    if os.path.exists(path):
                        pythonExe = path
                        logger.debug(f"Found Python executable at: {pythonExe}")
                        break
                
                if not os.path.exists(pythonExe):
                    logger.error(f"Python executable not found at expected path: {pythonExe}")
                    raise FileNotFoundError(f"Python executable not found in virtual environment")
            
            # 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 _executeCode(self, code: str, requirements: List[str] = None) -> Dict[str, Any]:
        """
        Execute Python code in the appropriate environment based on configuration.
        
        Args:
            code: Python code to execute
            requirements: List of required packages
            
        Returns:
            Execution result dictionary
        """
        # Check if we're in a production Azure environment
        env_type = APP_CONFIG.get("APP_ENV_TYPE", "dev").lower()
        
        logger.info(f"Executing code in environment type: {env_type}")
        
        if env_type == "prod":
            # Use the Azure-optimized execution method
            logger.info("Using Azure-optimized code execution method")
            return self._executeCodeProd(code, requirements)
        else:
            # Use the standard virtual environment execution method
            logger.info("Using standard virtual environment execution method")
            return self._executeCodeVenv(code, requirements)
        

    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()