pdd-cli 0.0.38__py3-none-any.whl → 0.0.40__py3-none-any.whl

pdd/crash_main.py

@@ -2,7 +2,7 @@ import sys
 from typing import Tuple, Optional, Dict, Any
 import click
 from rich import print as rprint
-from . import DEFAULT_STRENGTH
+from . import DEFAULT_STRENGTH, DEFAULT_TIME
 from pathlib import Path
 
 from .construct_paths import construct_paths
@@ -54,6 +54,11 @@ def crash_main(
     quiet = ctx.params.get("quiet", ctx.obj.get("quiet", False))
     verbose = ctx.params.get("verbose", ctx.obj.get("verbose", False)) # Get verbose flag
 
+    # Get model parameters from context early, including time
+    strength = ctx.obj.get("strength", DEFAULT_STRENGTH)
+    temperature = ctx.obj.get("temperature", 0)
+    time_param = ctx.obj.get("time", DEFAULT_TIME) # Renamed from time_budget for clarity
+
     try:
         # Construct file paths
         input_file_paths = {
@@ -84,71 +89,86 @@ def crash_main(
         program_content = input_strings["program_file"]
         error_content = input_strings["error_file"]
 
-        # Get model parameters from context
-        strength = ctx.obj.get("strength", DEFAULT_STRENGTH)
-        temperature = ctx.obj.get("temperature", 0)
+        # Store original content for comparison later
+        original_code_content = code_content
+        original_program_content = program_content
+
+        # Get model parameters from context (strength, temperature, time already fetched)
+        # strength = ctx.obj.get("strength", DEFAULT_STRENGTH) # Moved up
+        # temperature = ctx.obj.get("temperature", 0) # Moved up
+        # time_budget = ctx.obj.get("time", DEFAULT_TIME) # Moved up and renamed
 
         # verbose = ctx.params.get("verbose", ctx.obj.get("verbose", False)) # Already defined above
 
+        code_updated: bool = False
+        program_updated: bool = False
+
         if loop:
             # Use iterative fixing process
+            # Corrected parameter order for fix_code_loop, adding time_param
             success, final_program, final_code, attempts, cost, model = fix_code_loop(
-                code_file, prompt_content, program_file, strength, temperature, max_attempts or 3, budget or 5.0, error_file, verbose
+                code_file, prompt_content, program_file, strength, temperature,
+                max_attempts or 3, budget or 5.0, error_file, verbose, time_param
             )
+            # Determine if content was updated by fix_code_loop
+            if success: # Only consider updates if the loop reported success
+                code_updated = bool(final_code and final_code != original_code_content)
+                program_updated = bool(final_program and final_program != original_program_content)
         else:
             # Use single fix attempt
             if fix_code_module_errors is None:
                  raise ImportError("fix_code_module_errors is required but not available.")
-            # Note: fix_code_module_errors returns 7 values according to example
-            # update_program, update_code, fixed_program, fixed_code, program_code_fix, cost, model
-            # The current code unpacks 7 values, which matches the example.
-            update_program, update_code, final_program, final_code, program_code_fix, cost, model = fix_code_module_errors(
-                program_content, prompt_content, code_content, error_content, strength, temperature, verbose
+            # Corrected parameter order for fix_code_module_errors, adding time_param
+            fm_update_program, fm_update_code, final_program, final_code, program_code_fix, cost, model = fix_code_module_errors(
+                program_content, prompt_content, code_content, error_content, strength, temperature, verbose, time_param
             )
             success = True # Assume success after one attempt if no exception
             attempts = 1
+            # Use boolean flags from fix_code_module_errors and ensure content is not empty
+            code_updated = fm_update_code and bool(final_code)
+            program_updated = fm_update_program and bool(final_program)
 
-        # Ensure we have content to write, falling back to original content if needed
-        if final_program == "":
-            final_program = program_content
+        # Removed fallback to original content if final_code/final_program are empty
+        # An empty string from a fix function means no valid update.
 
-        if final_code == "":
-            final_code = code_content
-
-        # Determine whether to write the files based on whether paths are provided
+        # Determine whether to write the files based on whether paths are provided AND content was updated
         output_code_path_str = output_file_paths.get("output")
         output_program_path_str = output_file_paths.get("output_program")
 
-        should_write_code = output_code_path_str is not None
-        should_write_program = output_program_path_str is not None
-
-        # Write output files
-        if should_write_code:
+        # Write output files only if updated and path provided
+        if output_code_path_str and code_updated:
             output_code_path = Path(output_code_path_str)
             output_code_path.parent.mkdir(parents=True, exist_ok=True) # Ensure directory exists
             with open(output_code_path, "w") as f:
                 f.write(final_code)
 
-        if should_write_program:
+        if output_program_path_str and program_updated:
             output_program_path = Path(output_program_path_str)
             output_program_path.parent.mkdir(parents=True, exist_ok=True) # Ensure directory exists
             with open(output_program_path, "w") as f:
                 f.write(final_program)
 
-        # Provide user feedback (using quiet flag as per current implementation)
-        # To strictly follow the prompt's last note, change 'if not quiet:' to 'if verbose:'
+        # Provide user feedback
         if not quiet:
             if success:
-                rprint("[bold green]Crash fix completed successfully.[/bold green]")
+                rprint("[bold green]Crash fix attempt completed.[/bold green]") # Changed message slightly
             else:
-                rprint("[bold yellow]Crash fix completed with issues.[/bold yellow]")
+                rprint("[bold yellow]Crash fix attempt completed with issues or no changes made.[/bold yellow]") # Changed message
             rprint(f"[bold]Model used:[/bold] {model}")
             rprint(f"[bold]Total attempts:[/bold] {attempts}")
             rprint(f"[bold]Total cost:[/bold] ${cost:.2f}")
-            if should_write_code:
-                rprint(f"[bold]Fixed code saved to:[/bold] {output_code_path_str}")
-            if should_write_program:
-                rprint(f"[bold]Fixed program saved to:[/bold] {output_program_path_str}")
+
+            if output_code_path_str:
+                if code_updated:
+                    rprint(f"[bold]Fixed code saved to:[/bold] {output_code_path_str}")
+                else:
+                    rprint(f"[info]Code file {Path(code_file).name} was not modified. Output file {output_code_path_str} not written.[/info]")
+            
+            if output_program_path_str:
+                if program_updated:
+                    rprint(f"[bold]Fixed program saved to:[/bold] {output_program_path_str}")
+                else:
+                    rprint(f"[info]Program file {Path(program_file).name} was not modified. Output file {output_program_path_str} not written.[/info]")
 
         return success, final_code, final_program, attempts, cost, model
     

pdd/detect_change.py

@@ -1,9 +1,9 @@
-from typing import List, Dict, Tuple
+from typing import List, Dict, Tuple, Optional
 from pathlib import Path
 from rich.console import Console
 from rich.markdown import Markdown
 from pydantic import BaseModel, Field
-from . import EXTRACTION_STRENGTH
+from . import EXTRACTION_STRENGTH, DEFAULT_STRENGTH, DEFAULT_TIME
 
 from .preprocess import preprocess
 from .load_prompt_template import load_prompt_template
@@ -21,8 +21,9 @@ class ChangesList(BaseModel):
 def detect_change(
     prompt_files: List[str],
     change_description: str,
-    strength: float,
-    temperature: float,
+    strength: float=DEFAULT_STRENGTH,
+    temperature: float=0.0,
+    time: Optional[float] = DEFAULT_TIME,
     verbose: bool = False
 ) -> Tuple[List[Dict[str, str]], float, str]:
     """
@@ -33,6 +34,7 @@ def detect_change(
         change_description (str): Description of the changes to analyze
         strength (float): Strength parameter for the LLM model
         temperature (float): Temperature parameter for the LLM model
+        time (float): Time budget for LLM calls.
         verbose (bool): Whether to print detailed information
 
     Returns:
@@ -87,6 +89,7 @@ def detect_change(
             },
             strength=strength,
             temperature=temperature,
+            time=time,
             verbose=verbose
         )
 
@@ -104,6 +107,7 @@ def detect_change(
             input_json={"llm_output": detect_response['result']},
             strength=EXTRACTION_STRENGTH,
             temperature=0.0,
+            time=time,
             verbose=verbose,
             output_pydantic=ChangesList
         )

pdd/detect_change_main.py

@@ -6,6 +6,7 @@ from rich import print as rprint
 
 from .construct_paths import construct_paths
 from .detect_change import detect_change
+from . import DEFAULT_TIME
 
 def detect_change_main(
     ctx: click.Context,
@@ -55,6 +56,7 @@ def detect_change_main(
         # Get model parameters from context
         strength = ctx.obj.get('strength', 0.9)
         temperature = ctx.obj.get('temperature', 0)
+        time_budget = ctx.obj.get('time', DEFAULT_TIME)
 
         # Analyze which prompts need changes
         changes_list, total_cost, model_name = detect_change(
@@ -62,6 +64,7 @@ def detect_change_main(
             change_description,
             strength,
             temperature,
+            time_budget,
             verbose=not ctx.obj.get('quiet', False)
         )
 

pdd/fix_code_loop.py

@@ -4,6 +4,7 @@ import subprocess
 import sys
 from pathlib import Path
 from typing import Tuple
+from . import DEFAULT_TIME # Added DEFAULT_TIME
 
 # Use Rich for pretty printing to the console
 from rich.console import Console
@@ -25,6 +26,7 @@ except ImportError:
     # sys.path.append(str(Path(__file__).parent.parent)) # Adjust based on structure
     from fix_code_module_errors import fix_code_module_errors
 
+
 def fix_code_loop(
     code_file: str,
     prompt: str,
@@ -33,8 +35,9 @@ def fix_code_loop(
     temperature: float,
     max_attempts: int,
     budget: float,
-    error_log_file: str = "error_code.log",
+    error_log_file: str,
     verbose: bool = False,
+    time: float = DEFAULT_TIME,
 ) -> Tuple[bool, str, str, int, float, str | None]:
     """
     Attempts to fix errors in a code module through multiple iterations.
@@ -47,8 +50,9 @@ def fix_code_loop(
         temperature: LLM temperature (0.0 to 1.0).
         max_attempts: Maximum number of fix attempts.
         budget: Maximum cost allowed for the fixing process.
-        error_log_file: Path to the error log file (default: "error_code.log").
+        error_log_file: Path to the error log file.
         verbose: Enable detailed logging (default: False).
+        time: Time limit for the LLM calls (default: DEFAULT_TIME).
 
     Returns:
         Tuple containing the following in order:
@@ -219,6 +223,7 @@ def fix_code_loop(
                 errors=error_context_for_llm, # Pass the structured history
                 strength=strength,
                 temperature=temperature,
+                time=time, # Pass time
                 verbose=verbose,
             )
             if model_name_iter:

pdd/fix_code_module_errors.py

@@ -4,7 +4,7 @@ from rich import print
 from rich.markdown import Markdown
 from .load_prompt_template import load_prompt_template
 from .llm_invoke import llm_invoke
-from . import EXTRACTION_STRENGTH
+from . import EXTRACTION_STRENGTH, DEFAULT_TIME, DEFAULT_STRENGTH
 import json
 
 class CodeFix(BaseModel):
@@ -35,8 +35,9 @@ def fix_code_module_errors(
     prompt: str,
     code: str,
     errors: str,
-    strength: float,
+    strength: float = DEFAULT_STRENGTH,
     temperature: float = 0,
+    time: float = DEFAULT_TIME,
     verbose: bool = False
 ) -> Tuple[bool, bool, str, str, str, float, str]:
     """
@@ -72,6 +73,7 @@ def fix_code_module_errors(
             input_json=input_json,
             strength=strength,
             temperature=temperature,
+            time=time,
             verbose=verbose
         )
 
@@ -99,6 +101,7 @@ def fix_code_module_errors(
             input_json=extract_input,
             strength=EXTRACTION_STRENGTH,  # Fixed strength for extraction
             temperature=temperature,
+            time=time,
             verbose=verbose,
             output_pydantic=CodeFix
         )

pdd/fix_error_loop.py

@@ -11,6 +11,7 @@ from rich.console import Console
 
 # Relative import from an internal module.
 from .fix_errors_from_unit_tests import fix_errors_from_unit_tests
+from . import DEFAULT_TIME # Import DEFAULT_TIME
 
 console = Console()
 
@@ -104,7 +105,8 @@ def fix_error_loop(unit_test_file: str,
                    max_attempts: int,
                    budget: float,
                    error_log_file: str = "error_log.txt",
-                   verbose: bool = False):
+                   verbose: bool = False,
+                   time: float = DEFAULT_TIME):
     """
     Attempt to fix errors in a unit test and corresponding code using repeated iterations, 
     counting only the number of times we actually call the LLM fix function. 
@@ -124,6 +126,7 @@ def fix_error_loop(unit_test_file: str,
         budget: Maximum cost allowed for the fixing process.
         error_log_file: Path to file to log errors (default: "error_log.txt").
         verbose: Enable verbose logging (default: False).
+        time: Time parameter for the fix_errors_from_unit_tests call.
 
     Outputs:
         success: Boolean indicating if the overall process succeeded.
@@ -331,7 +334,8 @@ def fix_error_loop(unit_test_file: str,
                 error_log_file,
                 strength,
                 temperature,
-                verbose=verbose
+                verbose=verbose,
+                time=time # Pass time parameter
             )
             
             # Update the fix attempt in the structured log

pdd/fix_errors_from_unit_tests.py

@@ -10,6 +10,7 @@ from rich.panel import Panel
 from tempfile import NamedTemporaryFile
 
 from . import DEFAULT_STRENGTH
+from . import DEFAULT_TIME, EXTRACTION_STRENGTH
 from .preprocess import preprocess
 from .load_prompt_template import load_prompt_template
 from .llm_invoke import llm_invoke
@@ -104,9 +105,10 @@ def fix_errors_from_unit_tests(
     prompt: str,
     error: str,
     error_file: str,
-    strength: float,
-    temperature: float,
-    verbose: bool = False
+    strength: float = DEFAULT_STRENGTH,
+    temperature: float = 0.0,
+    verbose: bool = False,
+    time: float = DEFAULT_TIME
 ) -> Tuple[bool, bool, str, str, str, float, str]:
     """
     Fix errors in unit tests using LLM models and log the process.
@@ -120,6 +122,7 @@ def fix_errors_from_unit_tests(
         strength (float): LLM model strength (0-1)
         temperature (float): LLM temperature (0-1)
         verbose (bool): Whether to print detailed output
+        time (float): Time parameter for llm_invoke
 
     Returns:
         Tuple containing update flags, fixed code/tests, total cost, and model name
@@ -178,7 +181,8 @@ def fix_errors_from_unit_tests(
             },
             strength=strength,
             temperature=temperature,
-            verbose=verbose
+            verbose=verbose,
+            time=time
         )
 
         total_cost += response1['cost']
@@ -211,10 +215,11 @@ def fix_errors_from_unit_tests(
                 "unit_test": unit_test,
                 "code": code
             },
-            strength=DEFAULT_STRENGTH,  # Fixed strength as per requirements
+            strength=EXTRACTION_STRENGTH,  # Fixed strength as per requirements
             temperature=temperature,
             output_pydantic=CodeFix,
-            verbose=verbose
+            verbose=verbose,
+            time=time
         )
 
         total_cost += response2['cost']

pdd/fix_main.py

@@ -104,6 +104,8 @@ def fix_main(
         strength = ctx.obj.get('strength', DEFAULT_STRENGTH)
         temperature = ctx.obj.get('temperature', 0)
         verbose = ctx.obj.get('verbose', False)
+        time = ctx.obj.get('time') # Get time from context
+
         if loop:
             # Use fix_error_loop for iterative fixing
             success, fixed_unit_test, fixed_code, attempts, total_cost, model_name = fix_error_loop(
@@ -113,6 +115,7 @@ def fix_main(
                 verification_program=verification_program,
                 strength=strength,
                 temperature=temperature,
+                time=time, # Pass time to fix_error_loop
                 max_attempts=max_attempts,
                 budget=budget,
                 error_log_file=output_file_paths.get("output_results"),
@@ -128,6 +131,7 @@ def fix_main(
                 error_file=output_file_paths.get("output_results"),
                 strength=strength,
                 temperature=temperature,
+                time=time, # Pass time to fix_errors_from_unit_tests
                 verbose=verbose
             )
             success = update_unit_test or update_code

pdd/fix_verification_errors.py

@@ -5,6 +5,7 @@ from rich.markdown import Markdown
 from pydantic import BaseModel, Field
 from .load_prompt_template import load_prompt_template
 from .llm_invoke import llm_invoke
+from . import DEFAULT_TIME, DEFAULT_STRENGTH # Import defaults
 
 # Define Pydantic model for structured LLM output for VERIFICATION
 class VerificationOutput(BaseModel):
@@ -22,9 +23,10 @@ def fix_verification_errors(
     prompt: str,
     code: str,
     output: str,
-    strength: float,
+    strength: float = DEFAULT_STRENGTH,
     temperature: float = 0.0,
     verbose: bool = False,
+    time: float = DEFAULT_TIME
 ) -> Dict[str, Any]:
     """
     Identifies and fixes issues in a code module based on verification output.
@@ -37,6 +39,7 @@ def fix_verification_errors(
         strength: The strength (0-1) for the LLM model selection.
         temperature: The temperature for the LLM model. Defaults to 0.
         verbose: If True, prints detailed execution information. Defaults to False.
+        time: The time for the LLM model. Defaults to DEFAULT_TIME.
 
     Returns:
         A dictionary containing:
@@ -116,7 +119,8 @@ def fix_verification_errors(
             strength=strength,
             temperature=temperature,
             verbose=False,
-            output_pydantic=VerificationOutput
+            output_pydantic=VerificationOutput,
+            time=time
         )
         total_cost += verification_response.get('cost', 0.0)
         model_name = verification_response.get('model_name', model_name)
@@ -246,7 +250,8 @@ def fix_verification_errors(
                 strength=strength,
                 temperature=temperature,
                 verbose=False,
-                output_pydantic=FixerOutput
+                output_pydantic=FixerOutput,
+                time=time
             )
             total_cost += fix_response.get('cost', 0.0)
             model_name = fix_response.get('model_name', model_name)

pdd/fix_verification_errors_loop.py

@@ -2,10 +2,10 @@ import os
 import shutil
 import subprocess
 import datetime
-import time
 from pathlib import Path
 from typing import Dict, Tuple, Any, Optional
 from xml.sax.saxutils import escape
+import time
 
 from rich.console import Console
 
@@ -24,6 +24,8 @@ except ImportError:
             "Ensure it's available via relative import or in the 'pdd' package."
         )
 
+from . import DEFAULT_TIME # Import DEFAULT_TIME
+
 # Initialize Rich Console for pretty printing
 console = Console()
 
@@ -90,6 +92,7 @@ def fix_verification_errors_loop(
     output_program_path: Optional[str] = None,
     verbose: bool = False,
     program_args: Optional[list[str]] = None,
+    llm_time: float = DEFAULT_TIME # Add time parameter
 ) -> Dict[str, Any]:
     """
     Attempts to fix errors in a code file based on program execution output
@@ -109,6 +112,7 @@ def fix_verification_errors_loop(
         output_program_path: Optional path to save fixed program (default: None).
         verbose: Enable verbose logging (default: False).
         program_args: Optional list of command-line arguments for the program_file.
+        llm_time: Time parameter for fix_verification_errors calls (default: DEFAULT_TIME).
 
     Returns:
         A dictionary containing:
@@ -227,7 +231,8 @@ def fix_verification_errors_loop(
             output=initial_output,
             strength=strength,
             temperature=temperature,
-            verbose=verbose
+            verbose=verbose,
+            time=llm_time # Pass time
         )
         # 3e: Add cost
         initial_cost = initial_fix_result.get('total_cost', 0.0)
@@ -427,7 +432,8 @@ def fix_verification_errors_loop(
                 output=program_output,
                 strength=strength,
                 temperature=temperature,
-                verbose=verbose
+                verbose=verbose,
+                time=llm_time # Pass time
             )
 
             # 4f: Add cost

pdd/fix_verification_main.py

@@ -16,7 +16,7 @@ from .construct_paths import construct_paths
 from .fix_verification_errors import fix_verification_errors
 from .fix_verification_errors_loop import fix_verification_errors_loop
 # Import DEFAULT_STRENGTH from the main package
-from . import DEFAULT_STRENGTH
+from . import DEFAULT_STRENGTH, DEFAULT_TIME
 
 # Default values from the README
 DEFAULT_TEMPERATURE = 0.0
@@ -135,6 +135,7 @@ def fix_verification_main(
     force: bool = ctx.obj.get('force', False)
     quiet: bool = ctx.obj.get('quiet', False)
     verbose: bool = ctx.obj.get('verbose', False)
+    time: float = ctx.obj.get('time', DEFAULT_TIME) # Get time from context, default 0.25
 
     # --- Input Validation ---
     if loop and not verification_program:
@@ -250,35 +251,39 @@ def fix_verification_main(
     try:
         if loop:
             if not quiet:
-                rich_print("\n[bold blue]Running Iterative Verification (fix_verification_errors_loop)...[/bold blue]")
-
-            # fix_verification_errors_loop handles file I/O internally
-            # Pass file paths directly
-            # TODO: How are program_args passed? Assuming empty list for now. Needs CLI option.
-            program_args: List[str] = []
-
-            loop_results = fix_verification_errors_loop(
-                program_file=program_file,
-                code_file=code_file,
-                prompt=input_strings["prompt_file"],
-                verification_program=verification_program,
-                strength=strength,
-                temperature=temperature,
-                max_attempts=max_attempts,
-                budget=budget,
-                verification_log_file=output_results_path,
-                verbose=verbose,
-                program_args=program_args
-            )
-            success = loop_results['success']
-            final_program = loop_results['final_program']
-            final_code = loop_results['final_code']
-            attempts = loop_results['total_attempts']
-            total_cost = loop_results['total_cost']
-            model_name = loop_results['model_name']
-            # The loop function writes its own detailed log, so we don't need to build one here.
-            # We just need to ensure the path was passed correctly.
-
+                rich_print("[dim]Running Iterative Verification (fix_verification_errors_loop)...[/dim]")
+            try:
+                # Call fix_verification_errors_loop for iterative fixing
+                loop_results = fix_verification_errors_loop(
+                    program_file=program_file,                  # Changed to pass the program_file path
+                    code_file=code_file,                        # Changed to pass the code_file path
+                    prompt=input_strings["prompt_file"],        # Correctly passing prompt content
+                    verification_program=verification_program,      # Path to the verifier program
+                    strength=strength,
+                    temperature=temperature,
+                    llm_time=time, # Changed 'time' to 'llm_time'
+                    max_attempts=max_attempts,
+                    budget=budget,
+                    verification_log_file=output_results_path, # Use resolved output_results_path
+                    # output_code_path should not be passed here
+                    # output_program_path should not be passed here
+                    verbose=verbose,
+                    program_args=[] # Pass an empty list for program_args
+                )
+                success = loop_results.get('success', False)
+                final_program = loop_results.get('final_program', "") # Use .get for safety
+                final_code = loop_results.get('final_code', "")       # Use .get for safety
+                attempts = loop_results.get('total_attempts', 0)    # Use .get for safety
+                total_cost = loop_results.get('total_cost', 0.0)    # Use .get for safety
+                model_name = loop_results.get('model_name', "N/A")  # Use .get for safety
+                # Capture full statistics if available
+                # statistics = loop_results.get('statistics', {})
+            except Exception as e:
+                rich_print(f"[bold red]Error during loop execution:[/bold red] {e}")
+                if verbose:
+                    import traceback
+                    rich_print(Panel(traceback.format_exc(), title="Traceback", border_style="red"))
+                success = False
         else: # Single pass verification
             if not quiet:
                 rich_print("\n[bold blue]Running Single Pass Verification (fix_verification_errors)...[/bold blue]")
@@ -309,7 +314,8 @@ def fix_verification_main(
                 output=program_output,
                 strength=strength,
                 temperature=temperature,
-                verbose=verbose
+                verbose=verbose,
+                time=time # Pass time to single pass function
             )
 
             # Determine success: If no issues were found OR if fixes were applied

pdd/generate_test.py

@@ -2,7 +2,7 @@ from typing import Tuple, Optional
 from rich import print
 from rich.markdown import Markdown
 from rich.console import Console
-from . import EXTRACTION_STRENGTH
+from . import EXTRACTION_STRENGTH, DEFAULT_STRENGTH, DEFAULT_TIME
 from .load_prompt_template import load_prompt_template
 from .preprocess import preprocess
 from .llm_invoke import llm_invoke
@@ -15,9 +15,10 @@ console = Console()
 def generate_test(
     prompt: str,
     code: str,
-    strength: float,
-    temperature: float,
-    language: str,
+    strength: float=DEFAULT_STRENGTH,
+    temperature: float=0.0,
+    time: float = DEFAULT_TIME,
+    language: str = "python",
     verbose: bool = False
 ) -> Tuple[str, float, str]:
     """
@@ -29,6 +30,7 @@ def generate_test(
         strength (float): The strength of the LLM model (0-1).
         temperature (float): The temperature of the LLM model.
         language (str): The programming language for the unit test.
+        time (float, optional): Time budget for LLM calls. Defaults to DEFAULT_TIME.
         verbose (bool): Whether to print detailed information.
 
     Returns:
@@ -62,6 +64,7 @@ def generate_test(
             input_json=input_json,
             strength=strength,
             temperature=temperature,
+            time=time,
             verbose=verbose
         )
 
@@ -79,6 +82,7 @@ def generate_test(
             prompt_text=last_600_chars,
             strength=strength,
             temperature=temperature,
+            time=time,
             verbose=verbose
         )
         total_cost += check_cost
@@ -92,6 +96,7 @@ def generate_test(
                 llm_output=result,
                 strength=strength,
                 temperature=temperature,
+                time=time,
                 verbose=verbose
             )
             total_cost += continue_cost
@@ -104,6 +109,7 @@ def generate_test(
             language=language,
             strength=EXTRACTION_STRENGTH,
             temperature=temperature,
+            time=time,
             verbose=verbose
         )
         total_cost += post_cost