pdd-cli 0.0.2__py3-none-any.whl

pdd/fix_code_module_errors.py

@@ -0,0 +1,143 @@
+from typing import Tuple
+from pydantic import BaseModel, Field, ValidationError
+from rich import print
+from rich.markdown import Markdown
+from .load_prompt_template import load_prompt_template
+from .llm_invoke import llm_invoke
+import json
+
+class CodeFix(BaseModel):
+    update_program: bool = Field(description="Indicates if the program needs updating")
+    update_code: bool = Field(description="Indicates if the code module needs updating")
+    fixed_program: str = Field(description="The fixed program code")
+    fixed_code: str = Field(description="The fixed code module")
+
+def validate_inputs(
+    program: str,
+    prompt: str,
+    code: str,
+    errors: str,
+    strength: float
+) -> None:
+    """Validate input parameters."""
+    if not all([program, prompt, code, errors]):
+        raise ValueError("All string inputs (program, prompt, code, errors) must be non-empty")
+    
+    if not isinstance(strength, (int, float)):
+        raise ValueError("Strength must be a number")
+    
+    if not 0 <= strength <= 1:
+        raise ValueError("Strength must be between 0 and 1")
+
+def fix_code_module_errors(
+    program: str,
+    prompt: str,
+    code: str,
+    errors: str,
+    strength: float,
+    temperature: float = 0,
+    verbose: bool = False
+) -> Tuple[bool, bool, str, str, float, str]:
+    """
+    Fix errors in a code module that caused a program to crash and/or have errors.
+    """
+    try:
+        # Validate inputs
+        validate_inputs(program, prompt, code, errors, strength)
+
+        # Step 1: Load prompt templates
+        fix_prompt = load_prompt_template("fix_code_module_errors_LLM")
+        extract_prompt = load_prompt_template("extract_program_code_fix_LLM")
+        
+        if not all([fix_prompt, extract_prompt]):
+            raise ValueError("Failed to load one or more prompt templates")
+
+        total_cost = 0
+        model_name = ""
+
+        # Step 2: First LLM invoke for error analysis
+        input_json = {
+            "program": program,
+            "prompt": prompt,
+            "code": code,
+            "errors": errors
+        }
+
+        if verbose:
+            print("[blue]Running initial error analysis...[/blue]")
+
+        first_response = llm_invoke(
+            prompt=fix_prompt,
+            input_json=input_json,
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose
+        )
+
+        total_cost += first_response.get('cost', 0)
+        model_name = first_response.get('model_name', '')
+
+        if verbose:
+            print("[green]Error analysis complete[/green]")
+            print(Markdown(first_response['result']))
+            print(f"[yellow]Current cost: ${total_cost:.6f}[/yellow]")
+
+        # Step 4: Second LLM invoke for code extraction
+        extract_input = {
+            "program_code_fix": first_response['result'],
+            "program": program,
+            "code": code
+        }
+
+        if verbose:
+            print("[blue]Extracting code fixes...[/blue]")
+
+        second_response = llm_invoke(
+            prompt=extract_prompt,
+            input_json=extract_input,
+            strength=0.89,  # Fixed strength as specified
+            temperature=temperature,
+            verbose=verbose,
+            output_pydantic=CodeFix
+        )
+
+        total_cost += second_response.get('cost', 0)
+
+        # Step 5: Extract values from Pydantic result
+        result = second_response['result']
+
+        if isinstance(result, str):
+            try:
+                result_dict = json.loads(result)
+            except json.JSONDecodeError:
+                result_dict = {"result": result}
+            result = CodeFix.model_validate(result_dict)
+        elif isinstance(result, dict):
+            result = CodeFix.model_validate(result)
+        elif not isinstance(result, CodeFix):
+            result = CodeFix.model_validate({"result": str(result)})
+
+        if verbose:
+            print("[green]Code extraction complete[/green]")
+            print(f"[yellow]Total cost: ${total_cost:.6f}[/yellow]")
+            print(f"[blue]Model used: {model_name}[/blue]")
+
+        # Step 7: Return results
+        return (
+            result.update_program,
+            result.update_code,
+            result.fixed_program,
+            result.fixed_code,
+            total_cost,
+            model_name
+        )
+
+    except ValueError as ve:
+        print(f"[red]Value Error: {str(ve)}[/red]")
+        raise
+    except ValidationError:
+        print("[red]Validation Error: Invalid result format[/red]")
+        raise
+    except Exception as e:
+        print(f"[red]Unexpected error: {str(e)}[/red]")
+        raise

pdd/fix_error_loop.py

@@ -0,0 +1,216 @@
+import os
+import re
+import shutil
+import subprocess
+from dataclasses import dataclass
+from typing import Tuple, Optional
+from rich import print as rprint
+
+from .fix_errors_from_unit_tests import fix_errors_from_unit_tests
+
+@dataclass
+class IterationResult:
+    fails: int
+    errors: int
+    iteration: int
+    total_fails_and_errors: int
+
+    def is_better_than(self, other: Optional['IterationResult']) -> bool:
+        if other is None:
+            return True
+        if self.total_fails_and_errors < other.total_fails_and_errors:
+            return True
+        if self.total_fails_and_errors == other.total_fails_and_errors:
+            return self.errors < other.errors  # Prioritize fewer errors
+        return False
+
+def extract_test_results(pytest_output: str) -> Tuple[int, int]:
+    """Extract the number of fails and errors from pytest output.
+    
+    Args:
+        pytest_output (str): The complete pytest output text
+        
+    Returns:
+        Tuple[int, int]: Number of fails and errors respectively
+    """
+    fails = errors = 0
+    
+    # First try to match the summary line
+    summary_match = re.search(r'=+ (\d+) failed[\,\s]', pytest_output)
+    if summary_match:
+        fails = int(summary_match.group(1))
+    else:
+        # Fallback to looking for any "X failed" pattern
+        fail_match = re.search(r'(\d+)\s+failed', pytest_output)
+        if fail_match:
+            fails = int(fail_match.group(1))
+    
+    # Look for error patterns
+    error_match = re.search(r'(\d+)\s+error', pytest_output)
+    if error_match:
+        errors = int(error_match.group(1))
+        
+    return fails, errors
+
+def create_backup_files(unit_test_file: str, code_file: str, fails: int, 
+                       errors: int, iteration: int) -> Tuple[str, str]:
+    """Create backup files with iteration information in the filename."""
+    unit_test_backup = f"{os.path.splitext(unit_test_file)[0]}_{fails}_{errors}_{iteration}.py"
+    code_backup = f"{os.path.splitext(code_file)[0]}_{fails}_{errors}_{iteration}.py"
+    
+    shutil.copy2(unit_test_file, unit_test_backup)
+    shutil.copy2(code_file, code_backup)
+    
+    return unit_test_backup, code_backup
+
+def fix_error_loop(
+    unit_test_file: str,
+    code_file: str,
+    prompt: str,
+    verification_program: str,
+    strength: float,
+    temperature: float,
+    max_attempts: int,
+    budget: float,
+    error_log_file: str = "error_log.txt",
+    verbose: bool = False
+) -> Tuple[bool, str, str, int, float, str]:
+    """
+    Attempt to fix errors in a unit test and its corresponding code file through multiple iterations.
+    """
+    # Input validation
+    if not all([os.path.exists(f) for f in [unit_test_file, code_file, verification_program]]):
+        raise FileNotFoundError("One or more input files do not exist")
+    if not (0 <= strength <= 1 and 0 <= temperature <= 1):
+        raise ValueError("Strength and temperature must be between 0 and 1")
+    
+    # Step 1: Remove existing error log file if it exists
+    try:
+        if os.path.exists(error_log_file):
+            os.remove(error_log_file)
+    except FileNotFoundError:
+        pass  # File doesn't exist, which is fine
+    
+    # Step 2: Initialize variables
+    attempt_count = 0
+    total_cost = 0.0
+    best_iteration: Optional[IterationResult] = None
+    model_name = ""
+    
+    while attempt_count < max_attempts:
+        rprint(f"[bold yellow]Attempt {attempt_count + 1}[/bold yellow]")
+        
+        # Increment attempt counter first
+        attempt_count += 1
+        
+        # Step 3a: Run pytest
+        with open(error_log_file, 'a') as f:
+            result = subprocess.run(['python', '-m', 'pytest', '-vv', '--no-cov', unit_test_file],
+                                 capture_output=True, text=True)
+            f.write("\n****************************************************************************************************\n")
+            f.write("\nAttempt " + str(attempt_count) + ":\n")
+            f.write("\n****************************************************************************************************\n")
+            f.write(result.stdout + result.stderr)
+        
+        # Extract test results
+        fails, errors = extract_test_results(result.stdout)
+        current_iteration = IterationResult(fails, errors, attempt_count, fails + errors)
+        
+        # Step 3b: Check if tests pass
+        if fails == 0 and errors == 0:
+            break
+            
+        # Step 3c: Handle test failures
+        with open(error_log_file, 'r') as f:
+            error_content = f.read()
+        rprint(f"[bold red]Test output (attempt {attempt_count}):[/bold red]")
+        rprint(error_content.replace('[', '\\[').replace(']', '\\]'))
+        
+        # Create backups
+        backup_unit_test, backup_code = create_backup_files(
+            unit_test_file, code_file, fails, errors, attempt_count
+        )
+        
+        # Read current files
+        with open(unit_test_file, 'r') as f:
+            current_unit_test = f.read()
+        with open(code_file, 'r') as f:
+            current_code = f.read()
+            
+        # Try to fix errors
+        update_unit_test, update_code, fixed_unit_test, fixed_code, iteration_cost, model_name = (
+            fix_errors_from_unit_tests(
+                current_unit_test, current_code, prompt, error_content,
+                error_log_file, strength, temperature,
+                verbose=verbose
+            )
+        )
+        
+        total_cost += iteration_cost
+        if total_cost > budget:
+            rprint("[bold red]Budget exceeded![/bold red]")
+            break
+            
+        if not (update_unit_test or update_code):
+            rprint("[bold yellow]No changes needed or possible.[/bold yellow]")
+            break
+            
+        # Update files if needed
+        if update_unit_test:
+            with open(unit_test_file, 'w') as f:
+                f.write(fixed_unit_test)
+
+        if update_code:
+            with open(code_file, 'w') as f:
+                f.write(fixed_code)
+                
+            # Run verification
+            rprint("[bold yellow]Running Verification.[/bold yellow]")
+            verification_result = subprocess.run(['python', verification_program],
+                                              capture_output=True, text=True)
+
+            if verification_result.returncode != 0:
+                rprint("[bold red]Verification failed! Restoring previous code.[/bold red]")
+                shutil.copy2(backup_code, code_file)
+                with open(error_log_file, 'a') as f:
+                    f.write("****************************************************************************************************\n")
+                    f.write("\nVerification program failed! Here is the output and errors from the verification program that was running the code under test:\n" + verification_result.stdout + verification_result.stderr)
+                    f.write("****************************************************************************************************\n")
+                    f.write(f"\nRestoring previous working code.\n")
+                continue
+                
+        # Update best iteration if current is better
+        if current_iteration.is_better_than(best_iteration):
+            best_iteration = current_iteration
+            
+        # Check budget after increment
+        if total_cost > budget:
+            rprint("[bold red]Budget exceeded![/bold red]")
+            break
+            
+    # Step 4: Final test run
+    with open(error_log_file, 'a') as f:
+        final_result = subprocess.run(['python', '-m', 'pytest', '-vv', unit_test_file],
+                                    capture_output=True, text=True)
+        f.write("\nFinal test run:\n" + final_result.stdout + final_result.stderr)
+    rprint("[bold]Final test output:[/bold]")
+    rprint(final_result.stdout.replace('[', '\\[').replace(']', '\\]'))
+    
+    # Step 5: Restore best iteration if needed
+    final_fails, final_errors = extract_test_results(final_result.stdout)
+    if best_iteration and (final_fails + final_errors) > best_iteration.total_fails_and_errors:
+        rprint(f"[bold yellow]Restoring best iteration: {best_iteration.iteration} [/bold yellow]")
+        best_unit_test = f"{os.path.splitext(unit_test_file)[0]}_{best_iteration.fails}_{best_iteration.errors}_{best_iteration.iteration}.py"
+        best_code = f"{os.path.splitext(code_file)[0]}_{best_iteration.fails}_{best_iteration.errors}_{best_iteration.iteration}.py"
+        shutil.copy2(best_unit_test, unit_test_file)
+        shutil.copy2(best_code, code_file)
+    
+    # Step 6: Return results
+    with open(unit_test_file, 'r') as f:
+        final_unit_test = f.read()
+    with open(code_file, 'r') as f:
+        final_code = f.read()
+        
+    success = final_fails == 0 and final_errors == 0
+    
+    return success, final_unit_test, final_code, attempt_count, total_cost, model_name

pdd/fix_errors_from_unit_tests.py

@@ -0,0 +1,240 @@
+import os
+import tempfile  # Added missing import
+from datetime import datetime
+from typing import Tuple, Optional
+from pydantic import BaseModel, Field, ValidationError
+from rich import print as rprint
+from rich.markdown import Markdown
+from rich.console import Console
+from rich.panel import Panel
+from tempfile import NamedTemporaryFile
+
+from .preprocess import preprocess
+from .load_prompt_template import load_prompt_template
+from .llm_invoke import llm_invoke
+
+console = Console()
+
+class CodeFix(BaseModel):
+    update_unit_test: bool = Field(description="Whether the unit test needs to be updated")
+    update_code: bool = Field(description="Whether the code needs to be updated")
+    fixed_unit_test: str = Field(description="The fixed unit test code")
+    fixed_code: str = Field(description="The fixed code under test")
+
+def validate_inputs(strength: float, temperature: float) -> None:
+    """Validate strength and temperature parameters."""
+    if not 0 <= strength <= 1:
+        raise ValueError("Strength must be between 0 and 1")
+    if not 0 <= temperature <= 1:
+        raise ValueError("Temperature must be between 0 and 1")
+
+def write_to_error_file(file_path: str, content: str) -> None:
+    """Write content to error file with timestamp and separator."""
+    try:
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        separator = f"\n{'='*80}\n{timestamp}\n{'='*80}\n"
+        
+        # Ensure parent directory exists
+        parent_dir = os.path.dirname(file_path)
+        use_fallback = False
+        
+        if parent_dir:
+            try:
+                os.makedirs(parent_dir, exist_ok=True)
+            except Exception as e:
+                console.print(f"[yellow]Warning: Could not create directory {parent_dir}: {str(e)}[/yellow]")
+                # Fallback to system temp directory
+                use_fallback = True
+                parent_dir = None
+        
+        # Use atomic write with temporary file
+        try:
+            # First read existing content if file exists
+            existing_content = ""
+            if os.path.exists(file_path):
+                try:
+                    with open(file_path, 'r') as f:
+                        existing_content = f.read()
+                except Exception as e:
+                    console.print(f"[yellow]Warning: Could not read existing file {file_path}: {str(e)}[/yellow]")
+
+            # Write both existing and new content to temp file
+            with NamedTemporaryFile(mode='w', dir=parent_dir, delete=False) as tmp_file:
+                if existing_content:
+                    tmp_file.write(existing_content)
+                tmp_file.write(f"{separator}{content}\n")
+                tmp_path = tmp_file.name
+            
+            # Only attempt atomic move if not using fallback
+            if not use_fallback:
+                try:
+                    os.replace(tmp_path, file_path)
+                except Exception as e:
+                    console.print(f"[yellow]Warning: Could not move file to {file_path}: {str(e)}[/yellow]")
+                    use_fallback = True
+            
+            if use_fallback:
+                # Write to fallback location in system temp directory
+                fallback_path = os.path.join(tempfile.gettempdir(), os.path.basename(file_path))
+                try:
+                    os.replace(tmp_path, fallback_path)
+                    console.print(f"[yellow]Warning: Using fallback location: {fallback_path}[/yellow]")
+                except Exception as e:
+                    console.print(f"[red]Error writing to fallback location: {str(e)}[/red]")
+                    try:
+                        os.unlink(tmp_path)
+                    except:
+                        pass
+                    raise
+        except Exception as e:
+            console.print(f"[red]Error writing to error file: {str(e)}[/red]")
+            try:
+                os.unlink(tmp_path)
+            except:
+                pass
+            raise
+    except Exception as e:
+        console.print(f"[red]Error in write_to_error_file: {str(e)}[/red]")
+        raise
+
+def fix_errors_from_unit_tests(
+    unit_test: str,
+    code: str,
+    prompt: str,
+    error: str,
+    error_file: str,
+    strength: float,
+    temperature: float,
+    verbose: bool = False
+) -> Tuple[bool, bool, str, str, float, str]:
+    """
+    Fix errors in unit tests using LLM models and log the process.
+
+    Args:
+        unit_test (str): The unit test code
+        code (str): The code under test
+        prompt (str): The prompt that generated the code
+        error (str): The error message
+        error_file (str): Path to error log file
+        strength (float): LLM model strength (0-1)
+        temperature (float): LLM temperature (0-1)
+        verbose (bool): Whether to print detailed output
+
+    Returns:
+        Tuple containing update flags, fixed code/tests, total cost, and model name
+    """
+    # Input validation
+    if not all([unit_test, code, prompt, error, error_file]):
+        raise ValueError("All input parameters must be non-empty")
+    
+    validate_inputs(strength, temperature)
+
+    total_cost = 0.0
+    model_name = ""
+
+    try:
+        # Step 1: Load prompt templates
+        fix_errors_prompt = load_prompt_template("fix_errors_from_unit_tests_LLM")
+        extract_fix_prompt = load_prompt_template("extract_unit_code_fix_LLM")
+        
+        if not fix_errors_prompt or not extract_fix_prompt:
+            raise ValueError("Failed to load prompt templates")
+
+        # Step 2: Read error file content
+        existing_errors = ""
+        try:
+            if os.path.exists(error_file):
+                with open(error_file, 'r', encoding='utf-8') as f:
+                    existing_errors = f.read()
+        except Exception as e:
+            if verbose:
+                console.print(f"[yellow]Warning: Could not read error file: {str(e)}[/yellow]")
+
+        # Step 3: Run first prompt through llm_invoke
+        processed_prompt = preprocess(
+            prompt,
+            recursive=False,
+            double_curly_brackets=True,
+            exclude_keys=['unit_test', 'code', 'unit_test_fix']
+        )
+
+        if verbose:
+            console.print(Panel("[bold green]Running fix_errors_from_unit_tests...[/bold green]"))
+
+        response1 = llm_invoke(
+            prompt=fix_errors_prompt,
+            input_json={
+                "unit_test": unit_test,
+                "code": code,
+                "prompt": processed_prompt,
+                "errors": error
+            },
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose
+        )
+
+        total_cost += response1['cost']
+        model_name = response1['model_name']
+        result1 = response1['result']
+
+        # Step 4: Pretty print results and log to error file
+        if verbose:
+            console.print(Markdown(result1))
+            console.print(f"Cost of first run: ${response1['cost']:.6f}")
+
+        write_to_error_file(error_file, f"Model: {model_name}\nResult:\n{result1}")
+
+        # Step 5: Preprocess extract_fix prompt
+        processed_extract_prompt = preprocess(
+            extract_fix_prompt,
+            recursive=False,
+            double_curly_brackets=True,
+            exclude_keys=['unit_test', 'code', 'unit_test_fix']
+        )
+
+        # Step 6: Run second prompt through llm_invoke with fixed strength
+        if verbose:
+            console.print(Panel("[bold green]Running extract_unit_code_fix...[/bold green]"))
+
+        response2 = llm_invoke(
+            prompt=processed_extract_prompt,
+            input_json={
+                "unit_test_fix": result1,
+                "unit_test": unit_test,
+                "code": code
+            },
+            strength=0.895,  # Fixed strength as per requirements
+            temperature=temperature,
+            output_pydantic=CodeFix,
+            verbose=verbose
+        )
+
+        total_cost += response2['cost']
+        result2: CodeFix = response2['result']
+
+        if verbose:
+            console.print(f"Total cost: ${total_cost:.6f}")
+            console.print(f"Model used: {model_name}")
+
+        return (
+            result2.update_unit_test,
+            result2.update_code,
+            result2.fixed_unit_test,
+            result2.fixed_code,
+            total_cost,
+            model_name
+        )
+
+    except ValidationError as e:
+        error_msg = f"Validation error in fix_errors_from_unit_tests: {str(e)}"
+        if verbose:
+            console.print(f"[bold red]{error_msg}[/bold red]")
+        write_to_error_file(error_file, error_msg)
+        return False, False, "", "", 0.0, ""
+    except Exception as e:
+        error_msg = f"Error in fix_errors_from_unit_tests: {str(e)}"
+        if verbose:
+            console.print(f"[bold red]{error_msg}[/bold red]")
+        write_to_error_file(error_file, error_msg)
+        return False, False, "", "", 0.0, ""