main.py

import dspy
import subprocess
import tempfile
import os
import logging
import sys
import shutil
import time
from typing import Optional, Dict, Any
from dataclasses import dataclass
import json

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('code_generator.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

@dataclass
class Config:
    """Configuration class for the code generator."""
    llm_model: str = "ollama_chat/llama3"
    max_refinements: int = 3
    test_timeout: int = 30  # seconds
    pytest_verbose: bool = True
    temp_dir_prefix: str = "code_gen_"
    
    @classmethod
    def from_env(cls) -> 'Config':
        """Create config from environment variables."""
        return cls(
            llm_model=os.getenv('LLM_MODEL', cls.llm_model),
            max_refinements=int(os.getenv('MAX_REFINEMENTS', str(cls.max_refinements))),
            test_timeout=int(os.getenv('TEST_TIMEOUT', str(cls.test_timeout))),
            pytest_verbose=os.getenv('PYTEST_VERBOSE', 'true').lower() == 'true'
        )
    
    @classmethod
    def from_file(cls, config_path: str) -> 'Config':
        """Create config from JSON file."""
        try:
            with open(config_path, 'r') as f:
                config_data = json.load(f)
            return cls(**config_data)
        except FileNotFoundError:
            logger.warning(f"Config file {config_path} not found, using defaults")
            return cls()
        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in config file {config_path}: {e}")
            return cls()

def validate_environment() -> bool:
    """Validate that all required tools are available."""
    logger.info("Validating environment...")
    
    # Check if pytest is available
    if not shutil.which('pytest'):
        logger.error("pytest not found. Please install pytest: pip install pytest")
        return False
    
    # Check if ollama is available
    if not shutil.which('ollama'):
        logger.error("ollama not found. Please install Ollama from https://ollama.ai")
        return False
    
    # Check if ollama is running and has llama3
    try:
        result = subprocess.run(
            ['ollama', 'list'], 
            capture_output=True, 
            text=True, 
            timeout=10
        )
        if result.returncode != 0:
            logger.error("Ollama is not running. Please start Ollama: ollama serve")
            return False
        
        if 'llama3' not in result.stdout:
            logger.warning("Llama3 model not found. Attempting to pull...")
            pull_result = subprocess.run(
                ['ollama', 'pull', 'llama3'],
                capture_output=True,
                text=True,
                timeout=300  # 5 minutes for model download
            )
            if pull_result.returncode != 0:
                logger.error(f"Failed to pull Llama3 model: {pull_result.stderr}")
                return False
            logger.info("Successfully pulled Llama3 model")
            
    except subprocess.TimeoutExpired:
        logger.error("Timeout while checking Ollama status")
        return False
    except FileNotFoundError:
        logger.error("Ollama command not found")
        return False
    
    logger.info("Environment validation successful")
    return True

def validate_task_description(task_description: str) -> bool:
    """Validate the input task description."""
    if not task_description or not task_description.strip():
        logger.error("Task description cannot be empty")
        return False
    
    if len(task_description.strip()) < 10:
        logger.warning("Task description is very short, may lead to poor results")
    
    if len(task_description) > 1000:
        logger.warning("Task description is very long, may lead to poor results")
    
    return True

# --- 1. Configure DSPy ---
def configure_dspy(config: Config) -> bool:
    """Configure DSPy with error handling."""
    try:
        logger.info(f"Configuring DSPy with model: {config.llm_model}")
        llm = dspy.LM(config.llm_model)
        dspy.settings.configure(lm=llm)
        logger.info("DSPy configuration successful")
        return True
    except Exception as e:
        logger.error(f"Failed to configure DSPy: {e}")
        return False

# --- 2. Define DSPy Signatures ---

class CodeGeneratorSignature(dspy.Signature):
    """Generates Python code from a task description."""
    task_description = dspy.InputField(desc="A description of the programming task.")
    python_code = dspy.OutputField(desc="The generated Python code.")

class TestGeneratorSignature(dspy.Signature):
    """Generates unit tests for a given Python code."""
    python_code = dspy.InputField(desc="The Python code to test.")
    unit_test_code = dspy.OutputField(desc="The generated pytest unit tests.")

class CodeRefinerSignature(dspy.Signature):
    """Refines Python code based on test failures."""
    task_description = dspy.InputField(desc="The original programming task description.")
    failing_code = dspy.InputField(desc="The Python code that failed tests.")
    test_output = dspy.InputField(desc="The output from running tests (stdout/stderr).")
    refined_code = dspy.OutputField(desc="The refined Python code that should pass tests.")

# --- 3. Define the DSPy Program (Pipeline) ---

class CodeGenProgram(dspy.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.config = config
        self.code_generator = dspy.Predict(CodeGeneratorSignature)
        self.test_generator = dspy.Predict(TestGeneratorSignature)
        self.code_refiner = dspy.Predict(CodeRefinerSignature)
        logger.info("CodeGenProgram initialized successfully")

    def _execute_tests_safely(self, python_code: str, test_code: str) -> tuple[bool, str]:
        """Execute tests in a safe, isolated environment with proper error handling."""
        test_passed = False
        test_output = ""
        
        try:
            with tempfile.TemporaryDirectory(prefix=self.config.temp_dir_prefix) as temp_dir:
                test_file_path = os.path.join(temp_dir, "test_module.py")
                
                # Write code and tests to file with error handling
                try:
                    with open(test_file_path, "w", encoding='utf-8') as f:
                        f.write(python_code)
                        f.write("\n\n")
                        f.write(test_code)
                except (IOError, OSError) as e:
                    logger.error(f"Failed to write test file: {e}")
                    return False, f"File write error: {e}"

                # Build pytest command
                pytest_cmd = ["pytest"]
                if self.config.pytest_verbose:
                    pytest_cmd.append("-v")
                pytest_cmd.append(test_file_path)

                logger.debug(f"Executing: {' '.join(pytest_cmd)}")
                
                # Run pytest with timeout
                try:
                    process = subprocess.run(
                        pytest_cmd,
                        capture_output=True,
                        text=True,
                        cwd=temp_dir,
                        timeout=self.config.test_timeout
                    )
                    test_output = process.stdout + process.stderr
                    test_passed = (process.returncode == 0)
                    
                    if test_passed:
                        logger.info("Tests passed successfully")
                    else:
                        logger.warning(f"Tests failed with return code: {process.returncode}")
                        
                except subprocess.TimeoutExpired:
                    logger.error(f"Test execution timed out after {self.config.test_timeout} seconds")
                    test_output = f"Test execution timed out after {self.config.test_timeout} seconds"
                    test_passed = False
                except subprocess.SubprocessError as e:
                    logger.error(f"Subprocess error during test execution: {e}")
                    test_output = f"Subprocess error: {e}"
                    test_passed = False
                    
        except Exception as e:
            logger.error(f"Unexpected error during test execution: {e}")
            test_output = f"Unexpected error: {e}"
            test_passed = False
            
        return test_passed, test_output

    def _generate_code_safely(self, task_description: str) -> Optional[str]:
        """Generate code with error handling."""
        try:
            logger.info("Generating initial code...")
            code_prediction = self.code_generator(task_description=task_description)
            if hasattr(code_prediction, 'python_code') and code_prediction.python_code:
                logger.info("Code generation successful")
                return code_prediction.python_code
            else:
                logger.error("Code generation returned empty result")
                return None
        except Exception as e:
            logger.error(f"Error during code generation: {e}")
            return None

    def _generate_tests_safely(self, python_code: str) -> Optional[str]:
        """Generate tests with error handling."""
        try:
            logger.info("Generating tests...")
            test_prediction = self.test_generator(python_code=python_code)
            if hasattr(test_prediction, 'unit_test_code') and test_prediction.unit_test_code:
                logger.info("Test generation successful")
                return test_prediction.unit_test_code
            else:
                logger.error("Test generation returned empty result")
                return None
        except Exception as e:
            logger.error(f"Error during test generation: {e}")
            return None

    def _refine_code_safely(self, task_description: str, failing_code: str, test_output: str) -> Optional[str]:
        """Refine code with error handling."""
        try:
            logger.info("Refining code based on test failures...")
            refine_prediction = self.code_refiner(
                task_description=task_description,
                failing_code=failing_code,
                test_output=test_output
            )
            if hasattr(refine_prediction, 'refined_code') and refine_prediction.refined_code:
                logger.info("Code refinement successful")
                return refine_prediction.refined_code
            else:
                logger.error("Code refinement returned empty result")
                return None
        except Exception as e:
            logger.error(f"Error during code refinement: {e}")
            return None

    def forward(self, task_description: str, max_refinements: Optional[int] = None) -> dspy.Prediction:
        """Main forward method with comprehensive error handling and logging."""
        if max_refinements is None:
            max_refinements = self.config.max_refinements
            
        logger.info(f"Starting code generation pipeline for task: {task_description[:100]}...")
        
        # Validate input
        if not validate_task_description(task_description):
            return dspy.Prediction(
                python_code="",
                unit_test_code="",
                test_passed=False,
                final_test_output="Invalid task description",
                refinement_attempts=0,
                error="Invalid task description"
            )

        current_code = ""
        current_tests = ""
        test_passed = False
        test_output = ""
        refinement_attempts = 0

        for i in range(max_refinements + 1):  # +1 for initial generation
            logger.info(f"Iteration {i + 1}/{max_refinements + 1}")
            
            if i == 0:
                # Initial generation
                current_code = self._generate_code_safely(task_description)
                if not current_code:
                    logger.error("Failed to generate initial code")
                    return dspy.Prediction(
                        python_code="",
                        unit_test_code="",
                        test_passed=False,
                        final_test_output="Failed to generate initial code",
                        refinement_attempts=0,
                        error="Code generation failed"
                    )
                    
                current_tests = self._generate_tests_safely(current_code)
                if not current_tests:
                    logger.error("Failed to generate initial tests")
                    return dspy.Prediction(
                        python_code=current_code,
                        unit_test_code="",
                        test_passed=False,
                        final_test_output="Failed to generate tests",
                        refinement_attempts=0,
                        error="Test generation failed"
                    )
            else:
                # Refinement
                refinement_attempts += 1
                refined_code = self._refine_code_safely(task_description, current_code, test_output)
                if not refined_code:
                    logger.error(f"Failed to refine code on attempt {refinement_attempts}")
                    break
                    
                current_code = refined_code
                
                # Re-generate tests for refined code
                current_tests = self._generate_tests_safely(current_code)
                if not current_tests:
                    logger.error("Failed to generate tests for refined code")
                    break

            # Execute tests safely
            test_passed, test_output = self._execute_tests_safely(current_code, current_tests)

            if test_passed:
                logger.info(f"Success! Tests passed after {refinement_attempts} refinements")
                break
            else:
                logger.warning(f"Tests failed on iteration {i + 1}")

        # Log final results
        if test_passed:
            logger.info("Pipeline completed successfully")
        else:
            logger.warning(f"Pipeline completed without passing tests after {max_refinements} refinements")

        return dspy.Prediction(
            python_code=current_code,
            unit_test_code=current_tests,
            test_passed=test_passed,
            final_test_output=test_output,
            refinement_attempts=refinement_attempts,
            error=None if test_passed else "Tests did not pass after all refinement attempts"
        )

# --- 4. Define the Metric for Self-Correction ---

def test_execution_metric(example, prediction, trace=None) -> float:
    """
    Executes the generated tests against the generated code and returns 1.0 if they pass, 0.0 otherwise.
    Enhanced with proper error handling and logging.
    """
    if not hasattr(prediction, 'python_code') or not hasattr(prediction, 'unit_test_code'):
        logger.error("Prediction missing required attributes")
        return 0.0
        
    python_code = prediction.python_code
    unit_test_code = prediction.unit_test_code
    
    if not python_code or not unit_test_code:
        logger.error("Empty code or test content in prediction")
        return 0.0

    try:
        with tempfile.TemporaryDirectory(prefix="metric_test_") as temp_dir:
            test_file_path = os.path.join(temp_dir, "test_module.py")

            try:
                with open(test_file_path, "w", encoding='utf-8') as f:
                    f.write(python_code)
                    f.write("\n\n")
                    f.write(unit_test_code)
            except (IOError, OSError) as e:
                logger.error(f"Failed to write test file in metric: {e}")
                return 0.0

            # Run pytest using subprocess with timeout
            try:
                process = subprocess.run(
                    ["pytest", "-v", test_file_path],
                    capture_output=True,
                    text=True,
                    cwd=temp_dir,
                    timeout=30  # 30 second timeout for metric
                )
                
                # Return 1.0 for success (exit code 0), 0.0 for failure
                result = 1.0 if process.returncode == 0 else 0.0
                logger.debug(f"Metric test result: {result}")
                return result
                
            except subprocess.TimeoutExpired:
                logger.error("Metric test execution timed out")
                return 0.0
            except subprocess.SubprocessError as e:
                logger.error(f"Subprocess error in metric: {e}")
                return 0.0
                
    except Exception as e:
        logger.error(f"Unexpected error in test execution metric: {e}")
        return 0.0

# --- 5. Run the Program ---

def run_example_task(compiled_program: CodeGenProgram, config: Config) -> bool:
    """Run an example task to demonstrate the system."""
    # Define a new task to test the compiled program
    new_task = "Write a Python function that takes a list of integers and returns a new list containing only the even numbers, maintaining their original order."
    
    logger.info(f"Running example task: {new_task}")
    
    try:
        # Run the compiled program with the new task
        prediction = compiled_program(task_description=new_task, max_refinements=config.max_refinements)

        print("\n" + "="*60)
        print("GENERATED CODE")
        print("="*60)
        print(prediction.python_code)
        
        print("\n" + "="*60)
        print("GENERATED TESTS")
        print("="*60)
        print(prediction.unit_test_code)
        
        print("\n" + "="*60)
        print("EXECUTION RESULTS")
        print("="*60)
        print(f"Tests Passed: {prediction.test_passed}")
        print(f"Refinement Attempts: {getattr(prediction, 'refinement_attempts', 'Unknown')}")
        
        if hasattr(prediction, 'error') and prediction.error:
            print(f"Error: {prediction.error}")
            
        if not prediction.test_passed:
            print("\nFinal Test Output:")
            print("-" * 40)
            print(prediction.final_test_output)
            
        return prediction.test_passed
        
    except Exception as e:
        logger.error(f"Error running example task: {e}")
        return False

def main():
    """Main function with comprehensive error handling and configuration."""
    print("DSPy Self-Correcting Code Generator")
    print("====================================")
    
    # Load configuration
    config_path = os.getenv('CONFIG_PATH', 'config.json')
    if os.path.exists(config_path):
        logger.info(f"Loading configuration from {config_path}")
        config = Config.from_file(config_path)
    else:
        logger.info("Loading configuration from environment variables")
        config = Config.from_env()
    
    logger.info(f"Configuration: LLM={config.llm_model}, Max Refinements={config.max_refinements}")
    
    # Validate environment
    if not validate_environment():
        logger.error("Environment validation failed. Please check the requirements and try again.")
        sys.exit(1)
    
    # Configure DSPy
    if not configure_dspy(config):
        logger.error("DSPy configuration failed. Please check your LLM setup.")
        sys.exit(1)
    
    try:
        # Initialize the program
        logger.info("Initializing CodeGenProgram...")
        code_gen_program = CodeGenProgram(config)

        # Define a small training dataset
        train_data = [
            dspy.Example(
                task_description="Write a Python function that takes two integers and returns their sum.",
                python_code="def add_numbers(a, b):\n    return a + b",
                unit_test_code="import pytest\n\ndef test_add_numbers():\n    assert add_numbers(1, 2) == 3\ndef test_add_numbers_zero():\n    assert add_numbers(0, 0) == 0\ndef test_add_numbers_negative():\n    assert add_numbers(-1, -2) == -3"
            ).with_inputs("task_description"),
            dspy.Example(
                task_description="Write a Python function that takes a list of numbers and returns their product.",
                python_code="def product(numbers):\n    res = 1\n    for n in numbers:\n        res *= n\n    return res",
                unit_test_code="import pytest\n\ndef test_product_empty():\n    assert product([]) == 1\n\ndef test_product_single():\n    assert product([5]) == 5\n\ndef test_product_multiple():\n    assert product([1, 2, 3, 4]) == 24\n\ndef test_product_zero():\n    assert product([1, 2, 0, 4]) == 0"
            ).with_inputs("task_description")
        ]

        # Compile the program
        logger.info("Compiling the program with training data...")
        try:
            teleprompter = dspy.teleprompt.BootstrapFewShot(metric=test_execution_metric)
            compiled_program = teleprompter.compile(code_gen_program, trainset=train_data)
            logger.info("Program compilation successful")
        except Exception as e:
            logger.error(f"Program compilation failed: {e}")
            logger.info("Falling back to uncompiled program...")
            compiled_program = code_gen_program

        # Run example task
        success = run_example_task(compiled_program, config)
        
        if success:
            logger.info("Example task completed successfully!")
            print("\n✅ SUCCESS: Code generation and testing completed successfully!")
        else:
            logger.warning("Example task did not complete successfully")
            print("\n⚠️  WARNING: Code generation completed but tests did not pass")
            
    except KeyboardInterrupt:
        logger.info("Program interrupted by user")
        print("\n\nProgram interrupted by user.")
        sys.exit(0)
    except Exception as e:
        logger.error(f"Unexpected error in main: {e}")
        print(f"\n❌ ERROR: {e}")
        sys.exit(1)

def cli_interface():
    """Command-line interface for interactive usage."""
    config = Config.from_env()
    
    if not validate_environment():
        print("❌ Environment validation failed. Please check requirements.")
        return
        
    if not configure_dspy(config):
        print("❌ DSPy configuration failed.")
        return
        
    code_gen_program = CodeGenProgram(config)
    
    print("\nInteractive Code Generator")
    print("Type 'quit' to exit, 'help' for commands")
    
    while True:
        try:
            task = input("\nEnter task description: ").strip()
            
            if task.lower() in ['quit', 'exit', 'q']:
                break
            elif task.lower() == 'help':
                print("Commands:")
                print("  - Enter any task description to generate code")
                print("  - 'quit' or 'exit' to stop")
                print("  - 'help' to show this message")
                continue
            elif not task:
                continue
                
            prediction = code_gen_program(task_description=task)
            
            print(f"\n{'='*50}")
            print("GENERATED CODE:")
            print(prediction.python_code)
            print(f"\n{'='*50}")
            print("GENERATED TESTS:")
            print(prediction.unit_test_code)
            print(f"\nTests Passed: {prediction.test_passed}")
            
        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"Error: {e}")
    
    print("\nGoodbye!")

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == '--interactive':
        cli_interface()
    else:
        main()