pdd-cli 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

pdd/fix_errors_from_unit_tests.py

@@ -1,406 +1,248 @@
 import os
-import re
-import json
-import asyncio
-import tempfile
+import tempfile  # Added missing import
 from datetime import datetime
-from typing import Dict, Tuple, Any, Optional, List, Union
-import psutil  # Add psutil import for process management
-
-from rich.console import Console
+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 rich.text import Text
+from tempfile import NamedTemporaryFile
 
+from . import DEFAULT_STRENGTH
+from .preprocess import preprocess
 from .load_prompt_template import load_prompt_template
 from .llm_invoke import llm_invoke
-from .preprocess import preprocess
-from .edit_file import edit_file, run_edit_in_subprocess
-from langchain_mcp_adapters.client import MultiServerMCPClient
 
 console = Console()
 
-async def _fix_errors_from_unit_tests_async(
+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 = 0.0,
+    temperature: float,
     verbose: bool = False
 ) -> Tuple[bool, bool, str, str, str, float, str]:
     """
-    Fix unit test errors and warnings in code files.
-    
+    Fix errors in unit tests using LLM models and log the process.
+
     Args:
-        unit_test: The unit test code as a string
-        code: The code under test as a string
-        prompt: The prompt that generated the code under test
-        error: Errors and warnings that need to be fixed
-        error_file: Path to the file where error logs will be appended
-        strength: Strength of the LLM model to use (0-1)
-        temperature: Temperature for LLM output (0-1)
-        verbose: Whether to print detailed information
-        
+        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_unit_test: Boolean indicating if unit test was updated
-            - update_code: Boolean indicating if code was updated
-            - fixed_unit_test: The fixed unit test code
-            - fixed_code: The fixed code under test
-            - analysis_results: The raw output of the LLM analysis
-            - total_cost: Total cost of LLM invocations
-            - model_name: Name of the LLM model used
+        Tuple containing update flags, fixed code/tests, total cost, and model name
     """
-    # Initialize variables to track costs and model
+    # 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 = ""
-    
-    # Step 1: Load the prompt template
-    if verbose:
-        console.print("[bold blue]Step 1: Loading prompt template...[/bold blue]")
-    
-    console.print("[bold yellow]DEBUG: About to load prompt template[/bold yellow]")
-    prompt_template = load_prompt_template("fix_errors_from_unit_tests_LLM")
-    console.print(f"[bold yellow]DEBUG: Prompt template loaded: {'Success' if prompt_template else 'Failed'}[/bold yellow]")
-    
-    if not prompt_template:
-        error_msg = "Failed to load prompt template 'fix_errors_from_unit_tests_LLM'"
-        if verbose:
-            console.print(f"[bold red]{error_msg}[/bold red]")
-        raise ValueError(error_msg)
-        
-    if verbose:
-        console.print("[bold green]Prompt template loaded successfully[/bold green]")
-    
-    # Step 2: Read contents of error_file and parse any previous fix attempts
-    if verbose:
-        console.print("[bold blue]Step 2: Reading error file for previous fixes...[/bold blue]")
-    
-    prior_fixes = ""
+
     try:
-        if os.path.exists(error_file):
-            console.print("[bold yellow]DEBUG: Reading error file[/bold yellow]")
-            with open(error_file, 'r') as f:
-                prior_fixes = f.read()
-            
-            if verbose:
-                console.print(f"[bold green]Found existing error file: {error_file}[/bold green]")
-        else:
+        # 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"[bold yellow]Creating new error file: {error_file}[/bold yellow]")
-                
-            # Ensure directory exists
-            os.makedirs(os.path.dirname(os.path.abspath(error_file)), exist_ok=True)
-    except Exception as e:
-        if verbose:
-            console.print(f"[bold red]Error reading error file: {str(e)}[/bold red]")
-        prior_fixes = f"Error reading prior fixes: {str(e)}"
-    
-    # Step 3: Run the LLM analysis prompt through llm_invoke
-    if verbose:
-        console.print("[bold blue]Step 3: Running LLM analysis...[/bold blue]")
-    
-    # Preprocess the prompt
-    try:
-        console.print("[bold yellow]DEBUG: Preprocessing prompt[/bold yellow]")
+                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
+        )
+        
+        processed_fix_errors_prompt = preprocess(
+            fix_errors_prompt,
+            recursive=False,
             double_curly_brackets=True,
-            exclude_keys=['unit_test', 'code', 'unit_test_fix']
+            exclude_keys=['unit_test', 'code', 'errors', 'prompt']
         )
-        console.print("[bold yellow]DEBUG: Prompt preprocessed successfully[/bold yellow]")
-    except Exception as e:
-        processed_prompt = prompt
+
         if verbose:
-            console.print(f"[bold yellow]Error during prompt preprocessing, using original prompt: {str(e)}[/bold yellow]")
-    
-    # Prepare input for LLM
-    llm_input = {
-        'unit_test': unit_test,
-        'code': code,
-        'prompt': processed_prompt,
-        'errors': error,
-        'prior_fixes': prior_fixes
-    }
-    
-    # Log to console if verbose
-    if verbose:
-        console.print(Panel(
-            Text("Running LLM analysis", style="bold white"),
-            subtitle=f"Strength: {strength}, Temperature: {temperature}"
-        ))
-        console.print(f"Input tokens: {len(unit_test.split()) + len(code.split()) + len(processed_prompt.split()) + len(error.split())}")
-    
-    # Run the LLM analysis
-    try:
-        console.print("[bold yellow]DEBUG: About to invoke LLM[/bold yellow]")
-        llm_response = llm_invoke(
-            prompt=prompt_template,
-            input_json=llm_input,
+            console.print(Panel("[bold green]Running fix_errors_from_unit_tests...[/bold green]"))
+
+        response1 = llm_invoke(
+            prompt=processed_fix_errors_prompt,
+            input_json={
+                "unit_test": unit_test,
+                "code": code,
+                "prompt": processed_prompt,
+                "errors": error
+            },
             strength=strength,
             temperature=temperature,
             verbose=verbose
         )
-        console.print("[bold yellow]DEBUG: LLM invocation completed[/bold yellow]")
-        
-        # Update tracking variables
-        total_cost += llm_response['cost']
-        model_name = llm_response['model_name']
-        
-        # Extract response
-        analysis_results = llm_response['result']
-        
-        # Display response if verbose
-        if verbose:
-            console.print("\n[bold green]LLM Analysis Complete[/bold green]")
-            console.print(Markdown(analysis_results))
-            console.print(f"[bold]Output tokens: {llm_response.get('output_tokens', 'unknown')}[/bold]")
-            console.print(f"[bold]Cost: ${llm_response['cost']:.6f}[/bold]")
-    
-    except Exception as e:
-        error_msg = f"Error during LLM analysis: {str(e)}"
-        if verbose:
-            console.print(f"[bold red]{error_msg}[/bold red]")
-        
-        # Log the error to the error file
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        error_log = f"\n\n{'='*50}\nERROR LOG - {timestamp}\n{'='*50}\n{error_msg}\n"
-        
-        try:
-            with open(error_file, 'a') as f:
-                f.write(error_log)
-        except Exception as file_err:
-            if verbose:
-                console.print(f"[bold red]Failed to write to error file: {str(file_err)}[/bold red]")
-        
-        # Return default values
-        return False, False, unit_test, code, "", total_cost, model_name
-    
-    # Extract corrected code sections using regex
-    if verbose:
-        console.print("[bold blue]Step 3d: Extracting code sections...[/bold blue]")
-    
-    # Extract sections using regex
-    corrected_code_match = re.search(r'<corrected_code_under_test>(.*?)</corrected_code_under_test>', analysis_results, re.DOTALL)
-    corrected_unit_test_match = re.search(r'<corrected_unit_test>(.*?)</corrected_unit_test>', analysis_results, re.DOTALL)
-    
-    # Extract corrected code sections from the regex matches
-    corrected_code_text = ""
-    corrected_unit_test_text = ""
-    
-    if corrected_code_match:
-        corrected_code_text = corrected_code_match.group(1).strip()
-    
-    if corrected_unit_test_match:
-        corrected_unit_test_text = corrected_unit_test_match.group(1).strip()
-        
-    if verbose:
-        console.print(f"[bold yellow]Extracted code text: {bool(corrected_code_text)}[/bold yellow]")
-        console.print(f"[bold yellow]Extracted test text: {bool(corrected_unit_test_text)}[/bold yellow]")
-    
-    # Step 3c: Append the output to error_file
-    if verbose:
-        console.print("[bold blue]Step 3c: Logging analysis results...[/bold blue]")
-    
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    log_entry = f"\n\n{'='*50}\nANALYSIS LOG - {timestamp}\n{'='*50}\n{analysis_results}\n"
-    
-    try:
-        with open(error_file, 'a') as f:
-            f.write(log_entry)
-        
-        if verbose:
-            console.print(f"[bold green]Analysis logged to {error_file}[/bold green]")
-    except Exception as e:
-        if verbose:
-            console.print(f"[bold red]Failed to write to error file: {str(e)}[/bold red]")
-    
-    # Step 4: Pretty print the analysis results if verbose
-    if verbose:
-        console.print("[bold blue]Step 4: Displaying analysis results...[/bold blue]")
-        console.print(Panel(
-            Markdown(analysis_results),
-            title="Analysis Results",
-            expand=False
-        ))
-    
-    # Initialize variables for return values
-    update_unit_test = False
-    update_code = False
-    fixed_unit_test = unit_test
-    fixed_code = code
-    
-    # Step 5: Use edit_file to apply the fixes
-    if verbose:
-        console.print("[bold blue]Step 5: Applying fixes...[/bold blue]")
-    
-    # Step 5a: Apply unit test fixes if available
-    if corrected_unit_test_text:
-        try:
-            # Create a temporary file for the unit test
-            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
-                temp_test_file = temp_file.name
-                temp_file.write(unit_test)
-            
-            if verbose:
-                console.print(f"[bold]Applying unit test fixes...[/bold]")
-            
-            # Apply fixes using run_edit_in_subprocess for process isolation
-            test_success, test_error = run_edit_in_subprocess(
-                file_path=temp_test_file,
-                edit_instructions=corrected_unit_test_text
-            )
-            
-            # Read the modified file
-            if test_success and os.path.exists(temp_test_file):
-                with open(temp_test_file, 'r') as f:
-                    fixed_unit_test = f.read()
-                update_unit_test = True
-                
-                if verbose:
-                    console.print(f"[bold green]Unit test fixes applied successfully[/bold green]")
-            else:
-                if verbose:
-                    console.print(f"[bold red]Failed to apply unit test fixes: {test_error}[/bold red]")
-            
-            # Clean up
-            if os.path.exists(temp_test_file):
-                os.remove(temp_test_file)
-            
-        except Exception as e:
-            if verbose:
-                console.print(f"[bold red]Error applying unit test fixes: {str(e)}[/bold red]")
-    else:
+
+        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("[bold yellow]No unit test fixes required or provided[/bold yellow]")
-    
-    # Step 5b: Apply code fixes if available
-    if corrected_code_text:
-        try:
-            # Create a temporary file for the code
-            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
-                temp_code_file = temp_file.name
-                temp_file.write(code)
-            
-            if verbose:
-                console.print(f"[bold]Applying code fixes...[/bold]")
-            
-            # Apply fixes using run_edit_in_subprocess for process isolation
-            code_success, code_error = run_edit_in_subprocess(
-                file_path=temp_code_file,
-                edit_instructions=corrected_code_text
-            )
-            
-            # Read the modified file
-            if code_success and os.path.exists(temp_code_file):
-                with open(temp_code_file, 'r') as f:
-                    fixed_code = f.read()
-                update_code = True
-                
-                if verbose:
-                    console.print(f"[bold green]Code fixes applied successfully[/bold green]")
-            else:
-                if verbose:
-                    console.print(f"[bold red]Failed to apply code fixes: {code_error}[/bold red]")
-            
-            # Clean up
-            if os.path.exists(temp_code_file):
-                os.remove(temp_code_file)
-            
-        except Exception as e:
-            if verbose:
-                console.print(f"[bold red]Error applying code fixes: {str(e)}[/bold red]")
-    else:
+            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("[bold yellow]No code fixes required or provided[/bold yellow]")
-    
-    # Step 6: Return the results
-    if verbose:
-        console.print("[bold blue]Step 6: Returning results...[/bold blue]")
-        console.print(f"[bold green]Fix process completed[/bold green]")
-        console.print(f"[bold]Update unit test: {update_unit_test}[/bold]")
-        console.print(f"[bold]Update code: {update_code}[/bold]")
-        console.print(f"[bold]Total cost: ${total_cost:.6f}[/bold]")
-        console.print(f"[bold]Model used: {model_name}[/bold]")
-    
-    # One final cleanup of any lingering processes before returning
-    # terminate_mcp_processes()  # Removed as this function doesn't exist in edit_file.py
-    
-    return (
-        update_unit_test,
-        update_code,
-        fixed_unit_test,
-        fixed_code,
-        analysis_results,
-        total_cost,
-        model_name
-    )
+            console.print(Panel("[bold green]Running extract_unit_code_fix...[/bold green]"))
 
-def fix_errors_from_unit_tests(
-    unit_test: str,
-    code: str,
-    prompt: str,
-    error: str,
-    error_file: str,
-    strength: float,
-    temperature: float = 0.0,
-    verbose: bool = False
-) -> Tuple[bool, bool, str, str, str, float, str]:
-    """
-    Synchronous wrapper for fixing unit test errors and warnings in code files.
-    
-    Args:
-        unit_test: The unit test code as a string
-        code: The code under test as a string
-        prompt: The prompt that generated the code under test
-        error: Errors and warnings that need to be fixed
-        error_file: Path to the file where error logs will be appended
-        strength: Strength of the LLM model to use (0-1)
-        temperature: Temperature for LLM output (0-1)
-        verbose: Whether to print detailed information
-        
-    Returns:
-        Tuple containing:
-            - update_unit_test: Boolean indicating if unit test was updated
-            - update_code: Boolean indicating if code was updated
-            - fixed_unit_test: The fixed unit test code
-            - fixed_code: The fixed code under test
-            - analysis_results: The raw output of the LLM analysis
-            - total_cost: Total cost of LLM invocations
-            - model_name: Name of the LLM model used
-    """
-    # Input validation
-    if not isinstance(unit_test, str) or not isinstance(code, str) or not isinstance(prompt, str) or not isinstance(error, str):
-        raise ValueError("Input parameters must be strings")
-    
-    if not isinstance(error_file, str) or not error_file:
-        raise ValueError("error_file must be a non-empty string")
-    
-    if not isinstance(strength, float) or strength < 0 or strength > 1:
-        strength = max(0, min(strength, 1))  # Clamp to 0-1 range instead of raising error
-    
-    if not isinstance(temperature, float) or temperature < 0 or temperature > 1:
-        temperature = max(0, min(temperature, 1))  # Clamp to 0-1 range instead of raising error
-    
-    # Create and use new event loop instead of trying to get the current one (which causes deprecation warning)
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    
-    try:
-        # Run the async function and return results
-        return loop.run_until_complete(_fix_errors_from_unit_tests_async(
-            unit_test=unit_test,
-            code=code,
-            prompt=prompt,
-            error=error,
-            error_file=error_file,
-            strength=strength,
+        response2 = llm_invoke(
+            prompt=processed_extract_prompt,
+            input_json={
+                "unit_test_fix": result1,
+                "unit_test": unit_test,
+                "code": code
+            },
+            strength=DEFAULT_STRENGTH,  # Fixed strength as per requirements
             temperature=temperature,
+            output_pydantic=CodeFix,
             verbose=verbose
-        ))
-    finally:
-        # Clean up the loop
-        loop.close()
+        )
+
+        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,
+            result1,
+            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, ""