pdd-cli 0.0.2__py3-none-any.whl

pdd/git_update.py

@@ -0,0 +1,84 @@
+import os
+from typing import Tuple, Optional
+from rich import print
+from rich.console import Console
+from rich.panel import Panel
+from .update_prompt import update_prompt
+import git
+
+console = Console()
+
+def git_update(
+    input_prompt: str,
+    modified_code_file: str,
+    strength: float,
+    temperature: float,
+    verbose: bool = False
+) -> Tuple[Optional[str], float, str]:
+    """
+    Read in modified code, restore the prior checked-in version from GitHub,
+    update the prompt, write back the modified code, and return outputs.
+
+    Args:
+        input_prompt (str): The prompt that generated the original code.
+        modified_code_file (str): Filepath of the modified code.
+        strength (float): Strength parameter for the LLM model.
+        temperature (float): Temperature parameter for the LLM model.
+
+    Returns:
+        Tuple[Optional[str], float, str]: Modified prompt, total cost, and model name.
+    """
+    try:
+        # Check if inputs are valid
+        if not input_prompt or not modified_code_file:
+            raise ValueError("Input prompt and modified code file path are required.")
+
+        if not os.path.exists(modified_code_file):
+            raise FileNotFoundError(f"Modified code file not found: {modified_code_file}")
+
+        # Initialize git repository
+        repo = git.Repo(search_parent_directories=True)
+
+        # Get the file's relative path to the repo root
+        repo_root = repo.git.rev_parse("--show-toplevel")
+        relative_path = os.path.relpath(modified_code_file, repo_root)
+
+        # Read the modified code
+        with open(modified_code_file, 'r') as file:
+            modified_code = file.read()
+
+        # Restore the prior checked-in version
+        repo.git.checkout('HEAD', '--', relative_path)
+
+        # Read the original input code
+        with open(modified_code_file, 'r') as file:
+            original_input_code = file.read()
+
+        # Call update_prompt function
+        modified_prompt, total_cost, model_name = update_prompt(
+            input_prompt=input_prompt,
+            input_code=original_input_code,
+            modified_code=modified_code,
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose
+        )
+
+        # Write back the modified code
+        with open(modified_code_file, 'w') as file:
+            file.write(modified_code)
+
+
+        # Pretty print the results
+        console.print(Panel.fit(
+            f"[bold green]Success:[/bold green]\n"
+            f"Modified prompt: {modified_prompt}\n"
+            f"Total cost: ${total_cost:.6f}\n"
+            f"Model name: {model_name}"
+        ))
+
+        return modified_prompt, total_cost, model_name
+
+    except Exception as e:
+        console.print(Panel(f"[bold red]Error:[/bold red] {str(e)}", title="Error", expand=False))
+        return None, 0.0, ""

pdd/increase_tests.py

@@ -0,0 +1,93 @@
+from typing import Tuple
+from rich.console import Console
+
+from .load_prompt_template import load_prompt_template
+from .llm_invoke import llm_invoke
+from .postprocess import postprocess
+
+def increase_tests(
+    existing_unit_tests: str, 
+    coverage_report: str, 
+    code: str, 
+    prompt_that_generated_code: str,
+    language: str = "python",
+    strength: float = 0.5,
+    temperature: float = 0.0,
+    verbose: bool = False
+) -> Tuple[str, float, str]:
+    """
+    Generate additional unit tests to increase code coverage.
+
+    Args:
+        existing_unit_tests (str): Current unit tests for the code
+        coverage_report (str): Coverage report for the code
+        code (str): Code under test
+        prompt_that_generated_code (str): Original prompt used to generate the code
+        language (str, optional): Programming language. Defaults to "python".
+        strength (float, optional): LLM model strength. Defaults to 0.5.
+        temperature (float, optional): LLM model temperature. Defaults to 0.0.
+        verbose (bool, optional): Verbose output flag. Defaults to False.
+
+    Returns:
+        Tuple containing:
+        - Increased test function (str)
+        - Total cost of generation (float)
+        - Model name used (str)
+    """
+    console = Console()
+
+    # Validate inputs
+    if not all([existing_unit_tests, coverage_report, code, prompt_that_generated_code]):
+        raise ValueError("All input parameters must be non-empty strings")
+
+    # Validate strength and temperature
+    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")
+
+    try:
+        # Step 1: Load prompt template
+        prompt_name = "increase_tests_LLM"
+        prompt_template = load_prompt_template(prompt_name)
+        
+        if verbose:
+            console.print(f"[blue]Loaded Prompt Template:[/blue]\n{prompt_template}")
+
+        # Step 2: Prepare input for LLM invoke
+        input_json = {
+            "existing_unit_tests": existing_unit_tests,
+            "coverage_report": coverage_report,
+            "code": code,
+            "prompt_that_generated_code": prompt_that_generated_code,
+            "language": language
+        }
+
+        # Invoke LLM with the prompt
+        llm_response = llm_invoke(
+            prompt=prompt_template,
+            input_json=input_json,
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose
+        )
+
+        # Step 3: Postprocess the result
+        increase_test_function, total_cost, model_name = postprocess(
+            llm_response['result'], 
+            language, 
+            0.89,  # Same strength as LLM invoke
+            temperature, 
+            verbose
+        )
+
+        if verbose:
+            console.print(f"[green]Generated Test Function:[/green]\n{increase_test_function}")
+            console.print(f"[blue]Total Cost: ${total_cost:.6f}[/blue]")
+            console.print(f"[blue]Model Used: {model_name}[/blue]")
+
+        return increase_test_function, total_cost, model_name
+
+    except Exception as e:
+        console.print(f"[red]Error in increase_tests: {str(e)}[/red]")
+        raise

pdd/insert_includes.py

@@ -0,0 +1,150 @@
+from typing import Tuple
+from pathlib import Path
+from rich import print
+from pydantic import BaseModel, Field
+
+from .llm_invoke import llm_invoke
+from .load_prompt_template import load_prompt_template
+from .auto_include import auto_include
+from .preprocess import preprocess
+
+class InsertIncludesOutput(BaseModel):
+    output_prompt: str = Field(description="The prompt with dependencies inserted")
+
+def insert_includes(
+    input_prompt: str,
+    directory_path: str,
+    csv_filename: str,
+    strength: float,
+    temperature: float,
+    verbose: bool = False
+) -> Tuple[str, str, float, str]:
+    """
+    Determine needed dependencies and insert them into a prompt.
+
+    Args:
+        input_prompt (str): The prompt to process
+        directory_path (str): Directory path where the prompt file is located
+        csv_filename (str): Name of the CSV file containing dependencies
+        strength (float): Strength parameter for the LLM model
+        temperature (float): Temperature parameter for the LLM model
+        verbose (bool, optional): Whether to print detailed information. Defaults to False.
+
+    Returns:
+        Tuple[str, str, float, str]: Tuple containing:
+            - output_prompt: The prompt with dependencies inserted
+            - csv_output: Complete CSV output from auto_include
+            - total_cost: Total cost of running the function
+            - model_name: Name of the LLM model used
+    """
+    try:
+        # Step 1: Load the prompt template
+        insert_includes_prompt = load_prompt_template("insert_includes_LLM")
+        if not insert_includes_prompt:
+            raise ValueError("Failed to load insert_includes_LLM.prompt template")
+
+        if verbose:
+            print("[blue]Loaded insert_includes_LLM prompt template[/blue]")
+
+        # Step 2: Read the CSV file
+        try:
+            with open(csv_filename, 'r') as file:
+                csv_content = file.read()
+        except FileNotFoundError:
+            if verbose:
+                print(f"[yellow]CSV file {csv_filename} not found. Creating empty CSV.[/yellow]")
+            csv_content = "full_path,file_summary,date\n"
+            Path(csv_filename).write_text(csv_content)
+
+        # Step 3: Preprocess the prompt template
+        processed_prompt = preprocess(
+            insert_includes_prompt,
+            recursive=False,
+            double_curly_brackets=False
+        )
+
+        if verbose:
+            print("[blue]Preprocessed prompt template[/blue]")
+
+        # Step 4: Get dependencies using auto_include
+        dependencies, csv_output, auto_include_cost, auto_include_model = auto_include(
+            input_prompt=input_prompt,
+            directory_path=directory_path,
+            csv_file=csv_content,
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose
+        )
+
+        if verbose:
+            print("[blue]Retrieved dependencies using auto_include[/blue]")
+            print(f"Dependencies found: {dependencies}")
+
+        # Step 5: Run llm_invoke with the insert includes prompt
+        response = llm_invoke(
+            prompt=processed_prompt,
+            input_json={
+                "actual_prompt_to_update": input_prompt,
+                "actual_dependencies_to_insert": dependencies
+            },
+            strength=strength,
+            temperature=temperature,
+            verbose=verbose,
+            output_pydantic=InsertIncludesOutput
+        )
+
+        if not response or 'result' not in response:
+            raise ValueError("Failed to get valid response from LLM model")
+
+        result: InsertIncludesOutput = response['result']
+        model_name = response['model_name']
+        total_cost = response['cost'] + auto_include_cost
+
+        if verbose:
+            print("[green]Successfully inserted includes into prompt[/green]")
+            print(f"Total cost: ${total_cost:.6f}")
+            print(f"Model used: {model_name}")
+
+        return (
+            result.output_prompt,
+            csv_output,
+            total_cost,
+            model_name
+        )
+
+    except Exception as e:
+        print(f"[red]Error in insert_includes: {str(e)}[/red]")
+        raise
+
+def main():
+    """Example usage of the insert_includes function."""
+    # Example input
+    input_prompt = """% Generate a Python function that processes data
+    <include>data_processing.py</include>
+    """
+    directory_path = "./src"
+    csv_filename = "dependencies.csv"
+    strength = 0.7
+    temperature = 0.5
+
+    try:
+        output_prompt, csv_output, total_cost, model_name = insert_includes(
+            input_prompt=input_prompt,
+            directory_path=directory_path,
+            csv_filename=csv_filename,
+            strength=strength,
+            temperature=temperature,
+            verbose=True
+        )
+
+        print("\n[bold green]Results:[/bold green]")
+        print(f"[white]Output Prompt:[/white]\n{output_prompt}")
+        print(f"\n[white]CSV Output:[/white]\n{csv_output}")
+        print(f"[white]Total Cost: ${total_cost:.6f}[/white]")
+        print(f"[white]Model Used: {model_name}[/white]")
+
+    except Exception as e:
+        print(f"[red]Error in main: {str(e)}[/red]")
+
+if __name__ == "__main__":
+    main()

pdd/llm_invoke.py

@@ -0,0 +1,304 @@
+# llm_invoke.py
+
+import os
+import csv
+import json
+from pydantic import BaseModel, Field
+from rich import print as rprint
+
+from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
+from langchain_community.cache import SQLiteCache
+from langchain.globals import set_llm_cache
+from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
+from langchain_core.runnables import RunnablePassthrough, ConfigurableField
+
+from langchain_openai import AzureChatOpenAI
+from langchain_fireworks import Fireworks
+from langchain_anthropic import ChatAnthropic
+from langchain_openai import ChatOpenAI  # Chatbot and conversational tasks
+from langchain_openai import OpenAI  # General language tasks
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_together import Together
+from langchain_ollama.llms import OllamaLLM
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import LLMResult
+
+# import logging
+
+# Configure logging to output to the console
+# logging.basicConfig(level=logging.DEBUG)
+
+# Get the LangSmith logger
+# langsmith_logger = logging.getLogger("langsmith")
+
+# Set its logging level to DEBUG
+# langsmith_logger.setLevel(logging.DEBUG)
+
+class CompletionStatusHandler(BaseCallbackHandler):
+    def __init__(self):
+        self.is_complete = False
+        self.finish_reason = None
+        self.input_tokens = None
+        self.output_tokens = None
+
+    def on_llm_end(self, response: LLMResult, **kwargs) -> None:
+        self.is_complete = True
+        if response.generations and response.generations[0]:
+            generation = response.generations[0][0]
+            self.finish_reason = generation.generation_info.get('finish_reason', "").lower()
+
+            # Extract token usage
+            if hasattr(generation.message, 'usage_metadata'):
+                usage_metadata = generation.message.usage_metadata
+                self.input_tokens = usage_metadata.get('input_tokens')
+                self.output_tokens = usage_metadata.get('output_tokens')
+
+
+class ModelInfo:
+    def __init__(self, provider, model, input_cost, output_cost, coding_arena_elo,
+                 base_url, api_key, counter, encoder, max_tokens, max_completion_tokens,
+                 structured_output):
+        self.provider = provider.strip()
+        self.model = model.strip()
+        self.input_cost = float(input_cost) if input_cost else 0.0
+        self.output_cost = float(output_cost) if output_cost else 0.0
+        self.average_cost = (self.input_cost + self.output_cost) / 2
+        self.coding_arena_elo = float(coding_arena_elo) if coding_arena_elo else 0.0
+        self.base_url = base_url.strip() if base_url else None
+        self.api_key = api_key.strip() if api_key else None
+        self.counter = counter.strip() if counter else None
+        self.encoder = encoder.strip() if encoder else None
+        self.max_tokens = int(max_tokens) if max_tokens else None
+        self.max_completion_tokens = int(
+            max_completion_tokens) if max_completion_tokens else None
+        self.structured_output = structured_output.lower(
+        ) == 'true' if structured_output else False
+
+
+def load_models():
+    PDD_PATH = os.environ.get('PDD_PATH', '.')
+    # Assume that llm_model.csv is in PDD_PATH/data
+    models_file = os.path.join(PDD_PATH, 'data', 'llm_model.csv')
+    models = []
+    try:
+        with open(models_file, newline='') as csvfile:
+            reader = csv.DictReader(csvfile)
+            for row in reader:
+                model_info = ModelInfo(
+                    provider=row['provider'],
+                    model=row['model'],
+                    input_cost=row['input'],
+                    output_cost=row['output'],
+                    coding_arena_elo=row['coding_arena_elo'],
+                    base_url=row['base_url'],
+                    api_key=row['api_key'],
+                    counter=row['counter'],
+                    encoder=row['encoder'],
+                    max_tokens=row['max_tokens'],
+                    max_completion_tokens=row['max_completion_tokens'],
+                    structured_output=row['structured_output']
+                )
+                models.append(model_info)
+    except FileNotFoundError:
+        raise FileNotFoundError(f"llm_model.csv not found at {models_file}")
+    return models
+
+
+def select_model(strength, models, base_model_name):
+    # Get the base model
+    base_model = None
+    for model in models:
+        if model.model == base_model_name:
+            base_model = model
+            break
+    if not base_model:
+        raise ValueError(f"Base model {base_model_name} not found in the models list.")
+
+    if strength == 0.5:
+        return base_model
+    elif strength < 0.5:
+        # Models cheaper than or equal to the base model
+        cheaper_models = [
+            model for model in models if model.average_cost <= base_model.average_cost]
+        # Sort models by average_cost ascending
+        cheaper_models.sort(key=lambda m: m.average_cost)
+        if not cheaper_models:
+            return base_model
+        # Interpolate between cheapest model and base model
+        cheapest_model = cheaper_models[0]
+        cost_range = base_model.average_cost - cheapest_model.average_cost
+        target_cost = cheapest_model.average_cost + (strength / 0.5) * cost_range
+        # Find the model with closest average cost to target_cost
+        selected_model = min(
+            cheaper_models, key=lambda m: abs(m.average_cost - target_cost))
+        return selected_model
+    else:
+        # strength > 0.5
+        # Models better than or equal to the base model
+        better_models = [
+            model for model in models if model.coding_arena_elo >= base_model.coding_arena_elo]
+        # Sort models by coding_arena_elo ascending
+        better_models.sort(key=lambda m: m.coding_arena_elo)
+        if not better_models:
+            return base_model
+        # Interpolate between base model and highest ELO model
+        highest_elo_model = better_models[-1]
+        elo_range = highest_elo_model.coding_arena_elo - base_model.coding_arena_elo
+        target_elo = base_model.coding_arena_elo + \
+            ((strength - 0.5) / 0.5) * elo_range
+        # Find the model with closest ELO to target_elo
+        selected_model = min(
+            better_models, key=lambda m: abs(m.coding_arena_elo - target_elo))
+        return selected_model
+
+
+def create_llm_instance(selected_model, temperature, handler):
+    provider = selected_model.provider.lower()
+    model_name = selected_model.model
+    base_url = selected_model.base_url
+    api_key_name = selected_model.api_key
+    max_completion_tokens = selected_model.max_completion_tokens
+    max_tokens = selected_model.max_tokens
+
+    # Retrieve API key from environment variable if needed
+    api_key = os.environ.get(api_key_name) if api_key_name else None
+
+    # Initialize the appropriate LLM class
+    if provider == 'openai':
+        if base_url:
+            llm = ChatOpenAI(model=model_name, temperature=temperature,
+                                openai_api_key=api_key, callbacks=[handler], openai_api_base = base_url)
+        else:
+            if model_name[0] == 'o':
+                llm = ChatOpenAI(model=model_name, temperature=temperature,
+                    openai_api_key=api_key, callbacks=[handler], 
+                    model_kwargs = {'reasoning_effort':'high'})
+            else:
+                llm = ChatOpenAI(model=model_name, temperature=temperature,
+                                openai_api_key=api_key, callbacks=[handler])
+    elif provider == 'anthropic':
+        llm = ChatAnthropic(model=model_name, temperature=temperature,
+                            callbacks=[handler])
+    elif provider == 'google':
+        llm = ChatGoogleGenerativeAI(
+            model=model_name, temperature=temperature, callbacks=[handler])
+    elif provider == 'ollama':
+        llm =  OllamaLLM(
+            model=model_name, temperature=temperature, callbacks=[handler])
+    elif provider == 'azure':
+        llm = AzureChatOpenAI(
+            model=model_name, temperature=temperature, callbacks=[handler])
+    elif provider == 'fireworks':
+        llm = Fireworks(model=model_name, temperature=temperature,
+                        callbacks=[handler])
+    elif provider == 'together':
+        llm = Together(model=model_name, temperature=temperature,
+                       callbacks=[handler])
+    elif provider == 'groq':
+        llm = ChatGroq(model_name=model_name, temperature=temperature,
+                       callbacks=[handler])
+    else:
+        raise ValueError(f"Unsupported provider: {selected_model.provider}")
+    if max_completion_tokens:
+         llm.model_kwargs = {"max_completion_tokens" : max_completion_tokens}
+    else:
+        # Set max tokens if available
+        if max_tokens:
+            if provider == 'google':
+                llm.max_output_tokens = max_tokens
+            else:
+                llm.max_tokens = max_tokens
+    return llm
+
+
+def calculate_cost(handler, selected_model):
+    input_tokens = handler.input_tokens or 0
+    output_tokens = handler.output_tokens or 0
+    input_cost_per_million = selected_model.input_cost
+    output_cost_per_million = selected_model.output_cost
+    # Cost is (tokens / 1_000_000) * cost_per_million
+    total_cost = (input_tokens / 1_000_000) * input_cost_per_million + \
+        (output_tokens / 1_000_000) * output_cost_per_million
+    return total_cost
+
+
+def llm_invoke(prompt, input_json, strength, temperature, verbose=False, output_pydantic=None):
+    # Validate inputs
+    if not prompt:
+        raise ValueError("Prompt is required.")
+    if input_json is None:
+        raise ValueError("Input JSON is required.")
+    if not isinstance(input_json, dict):
+        raise ValueError("Input JSON must be a dictionary.")
+
+    # Set up cache
+    set_llm_cache(SQLiteCache(database_path=".langchain.db"))
+
+    # Get default model
+    base_model_name = os.environ.get('PDD_MODEL_DEFAULT', 'gpt-4o-mini')
+
+    # Load models
+    models = load_models()
+
+    # Select model
+    selected_model = select_model(strength, models, base_model_name)
+
+    # Create the prompt template
+    try:
+        prompt_template = PromptTemplate.from_template(prompt)
+    except Exception as e:
+        raise ValueError(f"Invalid prompt template: {str(e)}")
+
+    # Create a handler to capture token counts
+    handler = CompletionStatusHandler()
+
+    # Prepare LLM instance
+    llm = create_llm_instance(selected_model, temperature, handler)
+
+    # Handle structured output if output_pydantic is provided
+    if output_pydantic:
+        pydantic_model = output_pydantic
+        parser = PydanticOutputParser(pydantic_object=pydantic_model)
+        # Handle models that support structured output
+        if selected_model.structured_output:
+            llm = llm.with_structured_output(pydantic_model)
+            chain = prompt_template | llm
+        else:
+            # Use parser after the LLM
+            chain = prompt_template | llm | parser
+    else:
+        # Output is a string
+        chain = prompt_template | llm | StrOutputParser()
+
+    # Run the chain
+    try:
+        result = chain.invoke(input_json)
+    except Exception as e:
+        raise RuntimeError(f"Error during LLM invocation: {str(e)}")
+
+    # Calculate cost
+    cost = calculate_cost(handler, selected_model)
+
+    # If verbose, print information
+    if verbose:
+        rprint(f"Selected model: {selected_model.model}")
+        rprint(
+            f"Per input token cost: ${selected_model.input_cost} per million tokens")
+        rprint(
+            f"Per output token cost: ${selected_model.output_cost} per million tokens")
+        rprint(f"Number of input tokens: {handler.input_tokens}")
+        rprint(f"Number of output tokens: {handler.output_tokens}")
+        rprint(f"Cost of invoke run: ${cost}")
+        rprint(f"Strength used: {strength}")
+        rprint(f"Temperature used: {temperature}")
+        try:
+            rprint(f"Input JSON: {input_json}")
+        except:
+            print(f"Input JSON: {input_json}")
+        if output_pydantic:
+            rprint(f"Output Pydantic: {output_pydantic}")
+        rprint(f"Result: {result}")
+
+    return {'result': result, 'cost': cost, 'model_name': selected_model.model}

pdd/load_prompt_template.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+import os
+from typing import Optional
+import sys
+from rich import print
+
+def print_formatted(message: str) -> None:
+    """Print message with raw formatting tags for testing compatibility."""
+    print(message)
+
+def load_prompt_template(prompt_name: str) -> Optional[str]:
+    """
+    Load a prompt template from a file.
+
+    Args:
+        prompt_name (str): Name of the prompt file to load (without extension)
+
+    Returns:
+        str: The prompt template text
+    """
+    # Type checking
+    if not isinstance(prompt_name, str):
+        print_formatted("[red]Unexpected error loading prompt template[/red]")
+        return None
+
+    # Step 1: Get project path from environment variable
+    project_path = os.getenv('PDD_PATH')
+    if not project_path:
+        print_formatted("[red]PDD_PATH environment variable is not set[/red]")
+        return None
+
+    # Construct the full path to the prompt file
+    prompt_path = Path(project_path) / 'prompts' / f"{prompt_name}.prompt"
+
+    # Step 2: Load and return the prompt template
+    if not prompt_path.exists():
+        print_formatted(f"[red]Prompt file not found: {prompt_path}[/red]")
+        return None
+
+    try:
+        with open(prompt_path, 'r', encoding='utf-8') as file:
+            prompt_template = file.read()
+            print_formatted(f"[green]Successfully loaded prompt: {prompt_name}[/green]")
+            return prompt_template
+
+    except IOError as e:
+        print_formatted(f"[red]Error reading prompt file {prompt_name}: {str(e)}[/red]")
+        return None
+
+    except Exception as e:
+        print_formatted(f"[red]Unexpected error loading prompt template: {str(e)}[/red]")
+        return None
+
+if __name__ == "__main__":
+    # Example usage
+    prompt = load_prompt_template("example_prompt")
+    if prompt:
+        print_formatted("[blue]Loaded prompt template:[/blue]")
+        print_formatted(prompt)