pr-generator-agent 1.0.0__py3-none-any.whl

aipr/__init__.py

@@ -0,0 +1,18 @@
+"""
+AIPR - AI-powered Merge Request Description Generator
+"""
+
+import tomllib
+from pathlib import Path
+
+
+def get_version():
+    try:
+        pyproject_path = Path(__file__).parent.parent / "pyproject.toml"
+        with open(pyproject_path, "rb") as f:
+            return tomllib.load(f)["project"]["version"]
+    except Exception:
+        return "unknown"
+
+
+__version__ = get_version()

aipr/main.py

@@ -0,0 +1,634 @@
+import argparse
+import json
+import os
+import subprocess
+import sys
+from typing import Any, Dict, Optional, Tuple
+
+import git
+import tiktoken
+
+from .prompts import InvalidPromptError, PromptManager
+from .providers import generate_with_anthropic, generate_with_azure_openai, generate_with_openai
+
+# ANSI color codes
+BLUE = "\033[94m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+RED = "\033[91m"
+ENDC = "\033[0m"
+BOLD = "\033[1m"
+
+
+def detect_provider_and_model(model: Optional[str]) -> Tuple[str, str]:
+    """Detect which provider and model to use based on environment and args."""
+    if model:
+        # Handle simple aliases first
+        if model == "claude":
+            return "anthropic", "claude-3-sonnet-20240229"
+        if model == "azure":
+            return "azure", "gpt-4o-mini"
+        if model == "openai":
+            return "openai", "gpt-4o"
+
+        # Handle Azure models
+        if model.startswith("azure/"):
+            _, model_name = model.split("/", 1)
+            # Map azure model names to deployment names
+            azure_models = {
+                "o1-mini": "o1-mini",
+                "gpt-4o": "gpt-4o",
+                "gpt-4o-mini": "gpt-4o-mini",
+                "gpt-4": "gpt-4o",  # Alias for gpt-4o
+            }
+            return "azure", azure_models.get(model_name, model_name)
+
+        # Handle OpenAI models
+        if model.startswith("gpt") or model == "gpt4":
+            # Only use Azure if explicitly configured with endpoint
+            if os.getenv("AZURE_OPENAI_ENDPOINT") and model.startswith("azure/"):
+                openai_models = {
+                    "gpt4": "gpt-4",
+                    "gpt-4-turbo": "gpt-4-turbo",
+                    "gpt-4o": "gpt-4",
+                }
+                return "azure", openai_models.get(model, model)
+            else:
+                openai_models = {
+                    "gpt4": "gpt-4",
+                    "gpt-4-turbo": "gpt-4-turbo",
+                    "gpt-4o": "gpt-4",
+                }
+                return "openai", openai_models.get(model, model)
+
+        # Handle Anthropic models
+        if model.startswith("claude"):
+            anthropic_models = {
+                "claude-3": "claude-3-opus-20240229",
+                "claude-3-opus": "claude-3-opus-20240229",
+                "claude-3.5-sonnet": "claude-3-5-sonnet-20241022",
+                "claude-3-sonnet": "claude-3-sonnet-20240229",
+                "claude-3.5-haiku": "claude-3-5-haiku-20241022",
+                "claude-3-haiku": "claude-3-haiku-20240307",
+            }
+            return "anthropic", anthropic_models.get(model, model)
+
+    # No model specified, check environment for default
+    if os.getenv("ANTHROPIC_API_KEY"):
+        return "anthropic", "claude-3-sonnet-20240229"
+    if os.getenv("AZURE_OPENAI_ENDPOINT") and os.getenv("AZURE_API_KEY"):
+        return "azure", "gpt-4"
+    if os.getenv("OPENAI_API_KEY"):
+        return "openai", "gpt-4"
+
+    raise Exception(
+        "No API key found. Please set ANTHROPIC_API_KEY, AZURE_API_KEY, or OPENAI_API_KEY"
+    )
+
+
+def count_tokens(text: str, model: str) -> int:
+    """Count the number of tokens in a text string."""
+    try:
+        if model.startswith(("gpt-3", "gpt-4")):
+            encoding = tiktoken.encoding_for_model(model)
+        else:
+            # Default to cl100k_base for other models (including Azure and Claude)
+            encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(text))
+    except Exception:
+        # If we can't get a token count, return an estimate
+        return len(text) // 4  # Rough estimate of tokens
+
+
+def print_token_info(user_prompt: str, system_prompt: str, verbose: bool):
+    """Print token information for the prompts."""
+    if verbose:
+        print(f"System Prompt:\n{system_prompt}\n")
+        print(f"User Prompt:\n{user_prompt}\n")
+    # Add actual token counting if needed
+
+
+def print_separator(char="─", color=GREEN):
+    """Print a separator line with the given character and color."""
+    terminal_width = os.get_terminal_size().columns
+    print(f"{color}{char * terminal_width}{ENDC}")
+
+
+def print_header(text: str, level: int = 1):
+    """Print a header with the given text and level."""
+    if level == 1:
+        print(f"\n{text}")
+        print("=" * len(text))
+    else:
+        print(f"\n{text}")
+        print("-" * len(text))
+
+
+def run_trivy_scan(path: str, silent: bool = False, verbose: bool = False) -> Dict[str, Any]:
+    """Run trivy filesystem scan and return the results as a dictionary."""
+    try:
+        # Determine project type and scanning approach
+        is_java = os.path.exists(os.path.join(path, "pom.xml"))
+        is_node = os.path.exists(os.path.join(path, "package.json"))
+
+        # Enhanced Python project detection
+        python_files = [
+            "requirements.txt",
+            "setup.py",
+            "pyproject.toml",
+            "poetry.lock",
+            "Pipfile",
+            "Pipfile.lock",
+        ]
+        is_python = any(os.path.exists(os.path.join(path, f)) for f in python_files)
+
+        trivy_args = ["trivy", "fs", "--format", "json", "--scanners", "vuln,secret,config"]
+
+        # Add dependency scanning for supported project types
+        if is_java:
+            try:
+                if not silent:
+                    print(
+                        f"{BLUE}Detected Java project, " f"resolving Maven dependencies...{ENDC}",
+                        file=sys.stderr,
+                    )
+                subprocess.run(
+                    ["mvn", "dependency:resolve", "-DskipTests"],
+                    cwd=path,
+                    check=True,
+                    capture_output=True,
+                )
+                trivy_args.append("--dependency-tree")
+            except subprocess.CalledProcessError as e:
+                print(
+                    f"{YELLOW}Warning: Maven dependency resolution failed: {e}{ENDC}",
+                    file=sys.stderr,
+                )
+        elif is_node:
+            if not silent:
+                print(
+                    f"{BLUE}Detected Node.js project, including dependency scanning...{ENDC}",
+                    file=sys.stderr,
+                )
+            trivy_args.append("--dependency-tree")
+        elif is_python:
+            if not silent:
+                print(
+                    f"{BLUE}Detected Python project, including enhanced scanning...{ENDC}",
+                    file=sys.stderr,
+                )
+
+            # Check which package management files exist
+            pkg_files = []
+            has_poetry = os.path.exists(os.path.join(path, "poetry.lock"))
+            has_pipenv = os.path.exists(os.path.join(path, "Pipfile.lock"))
+            has_pip = os.path.exists(os.path.join(path, "requirements.txt"))
+            has_setup = os.path.exists(os.path.join(path, "setup.py"))
+            has_pyproject = os.path.exists(os.path.join(path, "pyproject.toml"))
+
+            if has_poetry:
+                pkg_files.append("poetry.lock")
+                if has_pyproject:
+                    pkg_files.append("pyproject.toml")
+                if not silent:
+                    print(
+                        f"{BLUE}Using Poetry for dependency scanning "
+                        f"(transitive dependencies, excludes dev)...{ENDC}",
+                        file=sys.stderr,
+                    )
+            elif has_pipenv:
+                pkg_files.append("Pipfile.lock")
+                if not silent:
+                    print(
+                        f"{BLUE}Using Pipenv for dependency scanning "
+                        f"(transitive dependencies, includes dev)...{ENDC}",
+                        file=sys.stderr,
+                    )
+            elif has_pip:
+                pkg_files.append("requirements.txt")
+                if not silent:
+                    print(
+                        f"{BLUE}Using pip requirements "
+                        f"(direct dependencies only, includes dev)...{ENDC}",
+                        file=sys.stderr,
+                    )
+            elif has_setup or has_pyproject:
+                if has_setup:
+                    pkg_files.append("setup.py")
+                if has_pyproject:
+                    pkg_files.append("pyproject.toml")
+                if not silent:
+                    print(f"{BLUE}Using Python package metadata files...{ENDC}", file=sys.stderr)
+
+            if pkg_files and not silent:
+                print(f"{BLUE}Found package files: {', '.join(pkg_files)}{ENDC}", file=sys.stderr)
+
+            # Add Python-specific scanning options
+            trivy_args.append("--dependency-tree")
+        else:
+            if not silent:
+                print(
+                    f"{BLUE}No specific package manager detected, "
+                    f"performing filesystem scan...{ENDC}",
+                    file=sys.stderr,
+                )
+
+        # Add the path as the last argument
+        trivy_args.append(path)
+
+        if not silent and verbose:
+            print(f"{BLUE}Running trivy with args: {' '.join(trivy_args)}{ENDC}", file=sys.stderr)
+
+        # Run trivy scan
+        result = subprocess.run(trivy_args, capture_output=True, text=True, check=True)
+        return json.loads(result.stdout)
+    except subprocess.CalledProcessError as e:
+        print(f"{RED}Error running trivy scan: {e}{ENDC}", file=sys.stderr)
+        if e.stderr:
+            print(f"{RED}Trivy error details: {e.stderr}{ENDC}", file=sys.stderr)
+        return {}
+    except json.JSONDecodeError as e:
+        print(f"{RED}Error parsing trivy output: {e}{ENDC}", file=sys.stderr)
+        return {}
+
+
+def compare_vulnerabilities(
+    current_scan: Dict[str, Any], target_scan: Dict[str, Any]
+) -> Tuple[str, str]:
+    """
+    Compare vulnerability scans between branches and return a tuple of
+    (markdown_report, analysis_data)
+    """
+    if not current_scan or not target_scan:
+        return "Error: Unable to generate vulnerability comparison", ""
+
+    report = ["## Vulnerability Comparison\n"]
+    analysis_data = ["### Security Analysis Data\n"]
+
+    # Get vulnerabilities from both scans
+    current_vulns = []
+    target_vulns = []
+
+    def extract_vulns(scan_data: Dict[str, Any]) -> list:
+        vulns = []
+        for result in scan_data.get("Results", []):
+            target = result.get("Target", "")
+            type = result.get("Type", "")
+            for vuln in result.get("Vulnerabilities", []):
+                vulns.append(
+                    {
+                        "id": vuln.get("VulnerabilityID"),
+                        "pkg": vuln.get("PkgName"),
+                        "version": vuln.get("InstalledVersion"),
+                        "severity": vuln.get("Severity"),
+                        "description": vuln.get("Description"),
+                        "fix_version": vuln.get("FixedVersion"),
+                        "target": target,
+                        "type": type,
+                        "title": vuln.get("Title"),
+                        "references": vuln.get("References", []),
+                    }
+                )
+        return vulns
+
+    current_vulns = extract_vulns(current_scan)
+    target_vulns = extract_vulns(target_scan)
+
+    # Create unique identifiers for comparison
+    def create_vuln_key(v: Dict[str, Any]) -> str:
+        return f"{v['id']}:{v['pkg']}:{v['version']}:{v['target']}"
+
+    current_vuln_keys = {create_vuln_key(v) for v in current_vulns}
+    target_vuln_keys = {create_vuln_key(v) for v in target_vulns}
+
+    # Find new and fixed vulnerabilities
+    new_vulns = current_vuln_keys - target_vuln_keys
+    fixed_vulns = target_vuln_keys - current_vuln_keys
+
+    # Group vulnerabilities by severity for analysis
+    severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3, "UNKNOWN": 4}
+
+    def group_by_severity(vulns_keys: set, vuln_list: list) -> Dict[str, list]:
+        grouped = {}
+        for vuln in vuln_list:
+            key = create_vuln_key(vuln)
+            if key in vulns_keys:
+                sev = vuln["severity"] or "UNKNOWN"
+                if sev not in grouped:
+                    grouped[sev] = []
+                grouped[sev].append(vuln)
+        return {
+            k: grouped[k] for k in sorted(grouped.keys(), key=lambda x: severity_order.get(x, 999))
+        }
+
+    # Prepare detailed analysis data
+    if new_vulns:
+        analysis_data.append("\nNew Vulnerabilities Details:")
+        grouped_new = group_by_severity(new_vulns, current_vulns)
+        for severity, vulns in grouped_new.items():
+            analysis_data.append(f"\n{severity} Severity:")
+            for vuln in sorted(vulns, key=lambda x: x["id"]):
+                analysis_data.append(f"\n- {vuln['id']} ({vuln['type']}):")
+                analysis_data.append(f"  - Package: {vuln['pkg']} {vuln['version']}")
+                analysis_data.append(f"  - In: {vuln['target']}")
+                analysis_data.append(f"  - Title: {vuln['title']}")
+                analysis_data.append(f"  - Description: {vuln['description']}")
+                if vuln["fix_version"]:
+                    fix_info = f"  - Fix available in version: " f"{vuln['fix_version']}"
+                    analysis_data.append(fix_info)
+                if vuln["references"]:
+                    analysis_data.append("  - References:")
+                    for ref in vuln["references"][:3]:  # Limit to first 3 references
+                        analysis_data.append(f"    * {ref}")
+
+    if fixed_vulns:
+        analysis_data.append("\nFixed Vulnerabilities Details:")
+        grouped_fixed = group_by_severity(fixed_vulns, target_vulns)
+        for severity, vulns in grouped_fixed.items():
+            analysis_data.append(f"\n{severity} Severity:")
+            for vuln in sorted(vulns, key=lambda x: x["id"]):
+                analysis_data.append(f"\n- {vuln['id']} ({vuln['type']}):")
+                analysis_data.append(f"  - Package: {vuln['pkg']} {vuln['version']}")
+                analysis_data.append(f"  - In: {vuln['target']}")
+                analysis_data.append(f"  - Title: {vuln['title']}")
+
+    # Generate markdown report
+    if new_vulns:
+        report.append("\n### New Vulnerabilities\n")
+        grouped_new = group_by_severity(new_vulns, current_vulns)
+        for severity, vulns in grouped_new.items():
+            report.append(f"\n#### {severity}\n")
+            for vuln in sorted(vulns, key=lambda x: x["id"]):
+                vuln_line = (
+                    f"- {vuln['id']} in {vuln['pkg']} " f"{vuln['version']} ({vuln['target']})"
+                )
+                report.append(vuln_line)
+
+    if fixed_vulns:
+        report.append("\n### Fixed Vulnerabilities\n")
+        grouped_fixed = group_by_severity(fixed_vulns, target_vulns)
+        for severity, vulns in grouped_fixed.items():
+            report.append(f"\n#### {severity}\n")
+            for vuln in sorted(vulns, key=lambda x: x["id"]):
+                vuln_line = (
+                    f"- {vuln['id']} in {vuln['pkg']} " f"{vuln['version']} ({vuln['target']})"
+                )
+                report.append(vuln_line)
+
+    if not new_vulns and not fixed_vulns:
+        report.append("\nNo vulnerability changes detected between branches.")
+        analysis_data.append("\nNo security changes to analyze.")
+
+    return "\n".join(report), "\n".join(analysis_data)
+
+
+class ColorHelpFormatter(argparse.RawDescriptionHelpFormatter):
+    """Custom formatter to preserve colors in help text."""
+
+    def _split_lines(self, text, width):
+        return text.splitlines()
+
+
+def parse_args(args=None):
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Generate PR description from git diff",
+        formatter_class=ColorHelpFormatter,
+        epilog=f"""
+recommended models:
+  {YELLOW}claude-3.5-sonnet{ENDC} (default)    Anthropic's Claude 3.5 Sonnet
+  {YELLOW}azure/o1-mini{ENDC}                  Azure OpenAI o1-mini
+  {YELLOW}azure/gpt-4o{ENDC}                   Azure OpenAI GPT-4
+  {GREEN}gpt-4{ENDC}                          OpenAI GPT-4
+
+prompt templates:
+  {BLUE}meta{ENDC}                          Default XML prompt template for merge requests""",
+    )
+    parser.add_argument("-s", "--silent", action="store_true", help="Silent mode")
+    parser.add_argument(
+        "-d",
+        "--debug",
+        action="store_true",
+        help="Debug mode - show prompts without sending to LLM",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="Verbose mode - show detailed API interaction",
+    )
+    parser.add_argument("-t", "--target", help="Target branch for comparison")
+    parser.add_argument("--vulns", action="store_true", help="Include vulnerability scan")
+    parser.add_argument("--working-tree", action="store_true", help="Use working tree")
+    parser.add_argument(
+        "-m",
+        "--model",
+        help="AI model to use (see recommended models below)",
+    )
+    parser.add_argument(
+        "-p",
+        "--prompt",
+        help=(
+            "Specify either a built-in prompt name (e.g., 'meta') or "
+            "a path to a custom XML prompt file (e.g., '~/prompts/custom.xml')"
+        ),
+    )
+    return parser.parse_args(args)
+
+
+def detect_default_branch(repo: git.Repo) -> str:
+    """Detect the default branch of the repository."""
+    for branch in ["main", "master", "develop"]:
+        try:
+            repo.git.rev_parse("--verify", branch)
+            return branch
+        except git.exc.GitCommandError:
+            continue
+    raise Exception("Could not detect default branch")
+
+
+def get_vulnerability_data() -> Optional[str]:
+    """Get vulnerability scan data using trivy."""
+    try:
+        result = subprocess.run(
+            ["trivy", "fs", "--quiet", "--severity", "HIGH,CRITICAL", "."],
+            capture_output=True,
+            text=True,
+        )
+        return result.stdout if result.stdout.strip() else None
+    except FileNotFoundError:
+        print("Warning: trivy not found. Skipping vulnerability scan.")
+        return None
+
+
+def generate_description(
+    diff: str,
+    vuln_data: Optional[str],
+    provider: str,
+    model: str,
+    system_prompt: str,
+    verbose: bool = False,
+    prompt_manager: Optional[PromptManager] = None,
+) -> str:
+    """Generate description using the specified provider."""
+    # Get the user prompt from the prompt manager
+    if prompt_manager is None:
+        prompt_manager = PromptManager()
+    user_prompt = prompt_manager.get_user_prompt(diff, vuln_data)
+
+    if provider == "anthropic":
+        return generate_with_anthropic(user_prompt, vuln_data, model, system_prompt, verbose)
+    if provider == "azure":
+        return generate_with_azure_openai(user_prompt, vuln_data, model, system_prompt, verbose)
+    if provider == "openai":
+        return generate_with_openai(user_prompt, vuln_data, model, system_prompt, verbose)
+    raise ValueError(f"Unknown provider: {provider}")
+
+
+def main(args=None):
+    """Main entry point for AIPR"""
+    args = parse_args(args)
+
+    try:
+        repo = git.Repo(os.getcwd(), search_parent_directories=True)
+    except git.InvalidGitRepositoryError:
+        print(f"{RED}Error: Directory is not a valid Git repository{ENDC}", file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        prompt_manager = PromptManager(args.prompt)
+    except InvalidPromptError as e:
+        print(f"{RED}Error: {str(e)}{ENDC}")
+        sys.exit(1)
+
+    provider, model = detect_provider_and_model(args.model)
+
+    # Get the diff based on the state
+    diff = ""
+    if args.target == "-" or repo.is_dirty():
+        # Show working tree changes
+        if not args.silent:
+            print(f"{BLUE}Showing working tree changes...{ENDC}", file=sys.stderr)
+        diff = repo.git.diff("HEAD", "--cached") + "\n" + repo.git.diff()
+    else:
+        # Compare with target branch
+        target = args.target
+        if not target:
+            # Try to find default branch
+            for branch in ["main", "master", "develop"]:
+                if branch in [h.name for h in repo.heads]:
+                    target = branch
+                    break
+
+        if target:
+            if not args.silent:
+                print(f"{BLUE}Comparing with {target}...{ENDC}", file=sys.stderr)
+            diff = repo.git.diff(f"{target}...{repo.active_branch.name}")
+        else:
+            print(f"{YELLOW}No suitable target branch found.{ENDC}", file=sys.stderr)
+            sys.exit(1)
+
+    if not diff.strip():
+        print("No changes found in the Git repository.", file=sys.stderr)
+        sys.exit(0)
+
+    # Get vulnerability data if requested
+    vuln_data = None
+    if args.vulns:
+        if not args.silent:
+            print(f"{BLUE}Running vulnerability scan...{ENDC}", file=sys.stderr)
+        vuln_data = run_trivy_scan(repo.working_dir, args.silent, False)
+
+    # Generate the description
+    if not args.silent:
+        print_header("\nGenerating Description")
+        print(f"Using {provider} ({model})...")
+
+    try:
+        # In debug mode, show the prompts that would be sent
+        if args.debug:
+            # Get the prompts first
+            system_prompt = prompt_manager.get_system_prompt()
+            user_prompt = prompt_manager.get_user_prompt(diff, vuln_data)
+
+            # Prepare the API parameters
+            if provider == "azure" and model == "o1-mini":
+                combined_prompt = (
+                    f"System Instructions:\n{system_prompt}\n\n" f"User Request:\n{user_prompt}"
+                )
+                messages = [{"role": "user", "content": combined_prompt}]
+                params = {
+                    "model": model,
+                    "messages": messages,
+                    "max_completion_tokens": 1000,
+                }
+            elif provider == "anthropic":
+                params = {
+                    "model": model,
+                    "system": system_prompt,
+                    "messages": [{"role": "user", "content": user_prompt}],
+                    "max_tokens": 1000,
+                    "temperature": 0.2,
+                }
+            else:
+                params = {
+                    "model": model,
+                    "messages": [
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": user_prompt},
+                    ],
+                    "max_tokens": 1000,
+                    "temperature": 0.2,
+                }
+
+            # Print debug information in a structured way
+            print_header("Debug Information")
+
+            print_header("API Call Structure", level=2)
+            print(f"Provider: {provider}")
+            print(f"Model: {model}")
+            print(
+                f"Endpoint: {os.getenv('AZURE_OPENAI_ENDPOINT', 'Not Set')}"
+                if provider == "azure"
+                else ""
+            )
+            print("\nParameters:")
+            print("─" * 40)
+            print(json.dumps({k: v for k, v in params.items() if k != "messages"}, indent=2))
+            print("\nMessages:")
+            print("─" * 40)
+            for msg in params["messages"]:
+                print(f"\n{msg['role'].upper()} MESSAGE:")
+                print(msg["content"])
+            print()
+            sys.exit(0)
+
+        # Generate the description
+        description = generate_description(
+            diff,
+            vuln_data,
+            provider,
+            model,
+            (
+                prompt_manager.get_system_prompt()
+                if prompt_manager
+                else PromptManager().get_system_prompt()
+            ),
+            args.verbose,
+            prompt_manager or PromptManager(),  # Ensure we always pass a valid PromptManager
+        )
+
+        if args.verbose:
+            print("\nAPI Response:")
+            print("─" * 40)
+        print(description)
+
+    except Exception as e:
+        print(f"{RED}Error: {provider.title()} API - {e}{ENDC}", file=sys.stderr)
+        sys.exit(1)
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

aipr/prompts/__init__.py

@@ -0,0 +1,5 @@
+"""Prompt templates for AIPR."""
+
+from .prompts import InvalidPromptError, PromptManager
+
+__all__ = ["PromptManager", "InvalidPromptError"]

aipr/prompts/meta.xml

@@ -0,0 +1,49 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<prompt>
+    <purpose>
+        Given a change-set and optional vulnerability-set generate a merge request description.
+        Strictly follow the instructions.
+        Format output like the example provided.
+        Always include a Security Analysis section.
+    </purpose>
+
+    <instructions>
+        <always>Focus only on the specific changes shown in the diff and vulnerability comparison.</always>
+        <always>Directly tie each point to actual code changes or security findings.</always>
+        <always>Highlight critical security changes when analyzing vulnerabilities.</always>
+        <always>Explain the impact of new vulnerabilities when provided.</always>
+        <always>Acknowledge any fixed vulnerabilities.</always>
+        <always>If tempted to write a concluding statement after the Security Analysis, do not continue.</always>
+        <never>Do not include generic concluding statements (e.g., "This improves the overall system").</never>
+        <never>Do not include broad claims about improvements (e.g., "This enhances development processes").</never>
+        <never>Do not include value judgments about the changes (e.g., "This is a significant improvement").</never>
+        <never>Do not include future benefits or implications.</never>
+    </instructions>
+
+    <example>
+        Change Summary:
+
+        1. **Added User Authentication**
+           - Implemented JWT middleware
+           - Added login/register endpoints
+           - Updated bcrypt to v5.1.1
+
+        2. **Security Updates**
+           - Fixed 2 medium severity vulnerabilities
+           - Updated deprecated crypto functions
+
+        Security Analysis:
+        ✓ No new vulnerabilities introduced
+        - Previous medium severity vulnerabilities in bcrypt have been resolved
+        - All crypto functions now use current best practices
+    </example>
+
+    <changes-set>
+        {{change_set}}
+    </changes-set>
+
+    <vulnerabilities-set>
+        {{vulnerability_set}}
+    </vulnerabilities-set>
+
+</prompt> 

aipr/prompts/prompts.py

@@ -0,0 +1,189 @@
+"""Prompt management for AIPR."""
+
+import importlib.resources
+import json
+import os
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Union
+
+
+class InvalidPromptError(Exception):
+    """Raised when a prompt file is invalid."""
+
+    pass
+
+
+class PromptManager:
+    """Manages system and user prompts for AI models."""
+
+    REQUIRED_XML_ELEMENTS = ["changes-set", "vulnerabilities-set"]
+
+    def __init__(self, prompt_name: str = None):
+        self._default_system_prompt = (
+            "You are a helpful assistant for generating Merge Requests.\n"
+            "Your task is to analyze Git changes and vulnerability comparison data to create "
+            "clear, well-structured merge request descriptions.\n"
+            "Response should end with the last specific change or security finding discussed.\n"
+            "If you find yourself wanting to write a concluding statement, stop writing instead."
+        )
+
+        self.prompt_name = prompt_name
+        self._xml_prompt = None
+
+        if prompt_name:
+            self._load_prompt(prompt_name)
+
+    def _load_prompt(self, prompt_name: str) -> None:
+        """Load a prompt from either a file path or a built-in prompt name."""
+        # First check if this is a file path (has .xml extension)
+        path = Path(os.path.expanduser(prompt_name))
+        if path.suffix == ".xml":
+            if path.exists():
+                self._load_xml_prompt(str(path))
+                return
+            # Try local prompts directory
+            local_path = Path("prompts") / path.name
+            if local_path.exists():
+                self._load_xml_prompt(str(local_path))
+                return
+            raise InvalidPromptError(f"Prompt file not found: {path}")
+
+        # If no .xml extension, treat as a built-in prompt name
+        available_prompts = self._get_available_prompts()
+        try:
+            with (
+                importlib.resources.files("aipr.prompts").joinpath(f"{prompt_name}.xml").open("r")
+            ) as f:
+                self._load_xml_prompt_from_string(f.read())
+        except Exception as e:
+            raise InvalidPromptError(
+                f"Could not load prompt '{prompt_name}'. "
+                f"Error: {e}\n\n"
+                f"Available built-in prompts: {', '.join(available_prompts)}"
+            )
+
+    def _validate_xml_prompt(self, root: ET.Element) -> None:
+        """Validate that the XML prompt has all required elements."""
+        for element in self.REQUIRED_XML_ELEMENTS:
+            if root.find(f".//{element}") is None:
+                raise InvalidPromptError(
+                    f"Invalid prompt file: Missing required element '{element}'"
+                )
+
+    def _load_xml_prompt(self, file_path: str) -> None:
+        """Load and parse the XML prompt template from a file."""
+        try:
+            tree = ET.parse(file_path)
+            root = tree.getroot()
+            self._validate_xml_prompt(root)
+            self._xml_prompt = root
+        except ET.ParseError as e:
+            raise InvalidPromptError(f"Error parsing XML prompt file: {e}")
+        except FileNotFoundError:
+            raise InvalidPromptError(f"Prompt file not found: {file_path}")
+        except PermissionError:
+            raise InvalidPromptError(f"Permission denied reading prompt file: {file_path}")
+
+    def _load_xml_prompt_from_string(self, xml_content: str) -> None:
+        """Load and parse the XML prompt template from a string."""
+        try:
+            root = ET.fromstring(xml_content)
+            self._validate_xml_prompt(root)
+            self._xml_prompt = root
+        except ET.ParseError as e:
+            raise InvalidPromptError(f"Error parsing XML prompt content: {e}")
+
+    def _get_available_prompts(self) -> list[str]:
+        """Get list of available built-in prompts."""
+        try:
+            prompts_dir = importlib.resources.files("aipr.prompts")
+            return [
+                f.stem for f in prompts_dir.iterdir() if f.suffix == ".xml" and f.stem != "__init__"
+            ]
+        except Exception:
+            return []
+
+    def get_system_prompt(self) -> str:
+        """Get the system prompt."""
+        return self._default_system_prompt
+
+    def get_user_prompt(self, diff: str, vuln_data: Union[str, dict, None] = None) -> str:
+        """Get the user prompt with diff and optional vulnerability data."""
+        if self.prompt_name:
+            if self._xml_prompt is None:
+                raise ValueError("XML prompt file was specified but could not be loaded")
+
+            # Find the changes-set and vulnerabilities-set elements
+            changes_set = self._xml_prompt.find(".//changes-set")
+            vulns_set = self._xml_prompt.find(".//vulnerabilities-set")
+
+            # Update the content
+            if changes_set is not None:
+                changes_set.text = "\n" + diff + "\n"  # Add newlines for better formatting
+            if vulns_set is not None:
+                if vuln_data:
+                    vulns_set.text = (
+                        "\n"
+                        + (
+                            json.dumps(vuln_data, indent=2)
+                            if isinstance(vuln_data, dict)
+                            else str(vuln_data)
+                        )
+                        + "\n"
+                    )
+                else:
+                    vulns_set.text = ""
+
+            # Get the example section
+            example = self._xml_prompt.find(".//example")
+            if example is not None:
+                example.tail = "\n"  # Add newline after example
+
+            # Convert to string while preserving formatting and remove XML declaration
+            xml_str = ET.tostring(self._xml_prompt, encoding="unicode", method="xml")
+            if xml_str.startswith("<?xml"):
+                xml_str = xml_str[xml_str.find("?>") + 2 :]
+
+            # Return the XML structure directly without wrapping it in additional text
+            return xml_str.strip()
+
+        # Default format when no XML prompt file is specified
+        prompt = [
+            "Please include:",
+            "- A concise summary of the changes",
+            "- Key modifications and their purpose",
+            "- Any notable technical details",
+            "- Security impact analysis (when vulnerability data is provided)",
+            "",
+            "Important Guidelines:",
+            "1. Focus only on the specific changes shown in the diff and vulnerability comparison",
+            "2. Each point must be directly tied to actual code changes or security findings",
+            "3. When analyzing vulnerabilities:",
+            "   - Highlight critical security changes",
+            "   - Explain the impact of new vulnerabilities",
+            "   - Acknowledge fixed vulnerabilities",
+            "4. DO NOT include any of the following:",
+            '   - Generic concluding statements (e.g., "This improves the overall system")',
+            '   - Broad claims about improvements (e.g., "This enhances development processes")',
+            '   - Value judgments about the changes (e.g., "This is a significant improvement")',
+            "   - Future benefits or implications",
+            "",
+            "Git Diff:",
+            diff,
+        ]
+
+        if vuln_data:
+            prompt.extend(
+                [
+                    "",
+                    "Vulnerability Analysis:",
+                    (
+                        json.dumps(vuln_data, indent=2)
+                        if isinstance(vuln_data, dict)
+                        else str(vuln_data)
+                    ),
+                ]
+            )
+
+        return "\n".join(prompt)

aipr/providers.py

@@ -0,0 +1,197 @@
+import json
+import os
+from typing import Any, Dict, Optional
+
+import anthropic
+from openai import AzureOpenAI, OpenAI
+
+
+def generate_with_anthropic(
+    diff: str,
+    vuln_data: Optional[Dict[str, Any]],
+    model: str,
+    system_prompt: str,
+    verbose: bool = False,
+) -> str:
+    """Generate description using Anthropic's Claude."""
+    if verbose:
+        print("\nInitializing Anthropic client...")
+
+    client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+
+    if verbose:
+        print("\nSending request to Anthropic API:")
+        print(f"  Model: {model}")
+        print("  Parameters:", json.dumps({"max_tokens": 1000, "temperature": 0.2}, indent=2))
+        print("\nRequest Messages:")
+        print("\nSYSTEM MESSAGE:")
+        print(system_prompt)
+        print("\nUSER MESSAGE:")
+        if len(diff) > 500:
+            print(diff[:500] + "...")
+        else:
+            print(diff)
+        print("\nMaking API call...")
+
+    try:
+        response = client.messages.create(
+            model=model,
+            max_tokens=1000,
+            temperature=0.2,
+            system=system_prompt,
+            messages=[{"role": "user", "content": diff}],
+        )
+        if verbose:
+            print("\nRaw API Response:")
+            print(f"  Model: {response.model}")
+            print(f"  Usage: {response.usage.model_dump() if response.usage else 'N/A'}")
+            print("\nResponse Content:")
+        return response.content[0].text
+    except Exception as e:
+        if verbose:
+            print(f"\nAPI Error: {str(e)}")
+        raise ValueError(f"Anthropic API error: {str(e)}")
+
+
+def generate_with_azure_openai(
+    diff: str,
+    vuln_data: Optional[Dict[str, Any]],
+    model: str,
+    system_prompt: str,
+    verbose: bool = False,
+) -> str:
+    """Generate description using Azure OpenAI."""
+    # Check required environment variables
+    endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+    api_key = os.getenv("AZURE_API_KEY")
+
+    if not endpoint or not api_key:
+        raise ValueError(
+            "Missing Azure OpenAI configuration. "
+            "Please set AZURE_OPENAI_ENDPOINT and AZURE_API_KEY environment variables."
+        )
+
+    try:
+        if verbose:
+            print("\nInitializing Azure OpenAI client with:")
+            print(f"  Endpoint: {endpoint}")
+            print("  API Version: 2024-02-15-preview")
+
+        client = AzureOpenAI(
+            api_key=api_key,
+            api_version="2024-02-15-preview",
+            azure_endpoint=endpoint,
+        )
+
+        # For models that don't support system messages (like o1-mini), prepend it to user message
+        messages = []
+        if model in ["o1-mini"]:
+            combined_prompt = f"System Instructions:\n{system_prompt}\n\n{diff}"
+            messages = [{"role": "user", "content": combined_prompt}]
+        else:
+            messages = [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": diff},
+            ]
+
+        # o1-mini has special parameter requirements:
+        # - doesn't support temperature parameter
+        kwargs = {
+            "model": model,
+            "messages": messages,
+        }
+        if model not in ["o1-mini"]:  # Only set these for non-o1-mini models
+            kwargs["max_tokens"] = 1000
+            kwargs["temperature"] = 0.2
+
+        if verbose:
+            print("\nSending request to Azure OpenAI API:")
+            print(f"  Model: {model}")
+            print(
+                "  Parameters:",
+                json.dumps({k: v for k, v in kwargs.items() if k != "messages"}, indent=2),
+            )
+            print("\nRequest Messages:")
+            for msg in messages:
+                print(f"\n{msg['role'].upper()} MESSAGE:")
+                content = msg["content"]
+                if len(content) > 500:
+                    print(content[:500] + "...")
+                else:
+                    print(content)
+
+        try:
+            if verbose:
+                print("\nMaking API call...")
+            response = client.chat.completions.create(**kwargs)
+            if not response.choices:
+                raise ValueError("No completion choices returned from the API")
+            if verbose:
+                print("\nRaw API Response:")
+                print(f"  Model: {response.model}")
+                print(f"  Usage: {response.usage.model_dump() if response.usage else 'N/A'}")
+                print("\nResponse Content:")
+            return response.choices[0].message.content
+        except Exception as api_error:
+            error_msg = str(api_error)
+            if verbose:
+                print(f"\nAPI Error: {error_msg}")
+            if "status" in error_msg:
+                raise ValueError(f"Azure OpenAI API error (HTTP {error_msg})")
+            raise ValueError(f"Azure OpenAI API error: {error_msg}")
+
+    except Exception as e:
+        if verbose:
+            print(f"\nProvider Error: {str(e)}")
+        raise ValueError(f"Azure OpenAI provider error: {str(e)}")
+
+
+def generate_with_openai(
+    diff: str,
+    vuln_data: Optional[Dict[str, Any]],
+    model: str,
+    system_prompt: str,
+    verbose: bool = False,
+) -> str:
+    """Generate description using OpenAI."""
+    if verbose:
+        print("\nInitializing OpenAI client...")
+
+    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": diff},
+    ]
+
+    if verbose:
+        print("\nSending request to OpenAI API:")
+        print(f"  Model: {model}")
+        print("  Parameters:", json.dumps({"max_tokens": 1000, "temperature": 0.2}, indent=2))
+        print("\nRequest Messages:")
+        for msg in messages:
+            print(f"\n{msg['role'].upper()} MESSAGE:")
+            content = msg["content"]
+            if len(content) > 500:
+                print(content[:500] + "...")
+            else:
+                print(content)
+        print("\nMaking API call...")
+
+    try:
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            max_tokens=1000,
+            temperature=0.2,
+        )
+        if verbose:
+            print("\nRaw API Response:")
+            print(f"  Model: {response.model}")
+            print(f"  Usage: {response.usage.model_dump() if response.usage else 'N/A'}")
+            print("\nResponse Content:")
+        return response.choices[0].message.content
+    except Exception as e:
+        if verbose:
+            print(f"\nAPI Error: {str(e)}")
+        raise ValueError(f"OpenAI API error: {str(e)}")

pr_generator_agent-1.0.0.dist-info/METADATA

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: pr-generator-agent
+Version: 1.0.0
+Summary: AI-powered Pull Request Description Generator
+Author-email: Daniel Scholl <daniel.scholl@microsoft.com>
+License: MIT
+License-File: LICENSE
+Requires-Python: >=3.11
+Requires-Dist: anthropic>=0.18.1
+Requires-Dist: gitpython>=3.1.42
+Requires-Dist: openai>=1.12.0
+Requires-Dist: tiktoken>=0.6.0
+Provides-Extra: dev
+Requires-Dist: black>=24.2.0; extra == 'dev'
+Requires-Dist: build>=1.0.3; extra == 'dev'
+Requires-Dist: flake8>=7.0.0; extra == 'dev'
+Requires-Dist: isort>=5.13.2; extra == 'dev'
+Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
+Requires-Dist: pytest>=8.0.2; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# AIPR - Agentic Pull Request Description Generator
+
+Automatically analyze git diffs and vulnerabilities to generate comprehensive, well-structured pull request descriptions. By intelligently detecting changes, performing security scans, and leveraging state-of-the-art AI models, AIPR helps teams save time while maintaining high-quality, consistent pull request descriptions.
+
+```bash
+# Install with pipx (recommended)
+pipx install pr-generator-agent
+
+# Or with pip
+pip install pr-generator-agent
+
+# Set the environment variable for the API key
+export ANTHROPIC_API_KEY="your-api-key"
+
+# Generate a PR description
+aipr
+
+# Custom usage - Analyze changes against main branch
+# Include: Vulnerability scanning
+# Use: Azure OpenAI o1-mini model
+# Prompt: meta template
+# Ouptut: Verbose
+aipr -t main --vulns -p meta -m azure/o1-mini -v
+
+# Inline with merge request creation
+gh pr create -b "$(aipr -s)" -t "feat: New Feature"
+```
+
+## Key Features
+
+- 🔍 **Smart Detection**: Automatically analyzes working tree changes or compares branches
+- 🛡️ **Security First**: Optional vulnerability scanning between branches using Trivy
+- 🤖 **AI-Powered**: Multiple AI providers (Azure OpenAI, OpenAI, Anthropic) for optimal results
+- 🔄 **CI/CD Ready**: Seamless integration with GitLab and GitHub workflows
+
+## Example Output
+
+```
+Change Summary:
+
+1. **Added User Authentication**
+   - Implemented JWT middleware
+   - Added login/register endpoints
+   - Updated bcrypt to v5.1.1
+
+2. **Security Updates**
+   - Fixed 2 medium severity vulnerabilities
+   - Updated deprecated crypto functions
+
+Security Analysis:
+✓ No new vulnerabilities introduced
+```
+
+## Requirements
+
+- Python 3.10 or higher (3.10, 3.11 officially supported)
+- Git
+- LLM API Key (Anthropic, OpenAI, or Azure OpenAI)
+- [Trivy](https://aquasecurity.github.io/trivy/latest/getting-started/installation/) (used for `--vulns` scanning)
+
+## Environment Variables
+
+#### Anthropic (Default)
+- `ANTHROPIC_API_KEY`: Anthropic API key
+
+#### Azure OpenAI
+- `AZURE_API_KEY`: Azure OpenAI API key
+- `AZURE_API_BASE`: Azure endpoint URL
+- `AZURE_API_VERSION`: API version (default: "2024-02-15-preview")
+
+#### OpenAI
+- `OPENAI_API_KEY`: OpenAI API key
+
+## Usage
+
+### Command Options
+- `-t, --target`: Compare changes with specific branch (default: auto-detects main/master)
+- `-p, --prompt`: Select prompt template (e.g., 'meta')
+- `-v, --verbose`: Show API interaction details
+- `-d, --debug`: Preview prompts without API calls
+- `-s, --silent`: Output only the description
+- `--vulns`: Include vulnerability scanning
+- `-m, --model`: Specify AI model to use
+
+The tool intelligently detects changes by:
+1. Using staged/unstaged changes if present
+2. Comparing against target branch if working tree is clean
+
+## Supported AI Models
+
+Choose from multiple AI providers:
+
+| Provider | Model | Notes |
+|----------|--------|-------|
+| Anthropic | `claude-3-sonnet` | default |
+| | `claude-3.5-sonnet` | latest |
+| | `claude-3.5-haiku` | latest |
+| | `claude-3-opus` | latest |
+| | `claude-3-haiku` | |
+| | `claude` | alias for `claude-3-sonnet` |
+| Azure OpenAI | `azure/o1-mini` | |
+| | `azure/gpt-4o-mini` | |
+| | `azure/gpt-4o` | |
+| | `azure` | alias for `azure/gpt-4o-mini` |
+| OpenAI | `gpt-4o` | |
+| | `gpt-4-turbo` | |
+| | `gpt-3.5-turbo` | |
+| | `openai` | alias for `gpt-4o` |
+
+## Custom Prompts
+
+AIPR supports custom prompt templates that allow you to tailor merge request descriptions to your team's specific needs. Custom prompts enable you to:
+- Define consistent formatting across your team
+- Include organization-specific requirements
+- Add custom sections and validation rules
+- Provide examples that match your team's standards
+
+For detailed information on creating and using custom prompts, see our [Custom Prompts Tutorial](docs/custom_prompts.md).
+
+## Contributing
+
+We welcome contributions! See our [Contributing Guide](CONTRIBUTING.md) for details on:
+- Setting up your development environment
+- Our development workflow
+- Code style guidelines
+- Pull request process
+- Running tests
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

pr_generator_agent-1.0.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+aipr/__init__.py,sha256=PW5IiHykwROknlAhWnyBPW5oPI01M3SwgKXrRX_oAv0,386
+aipr/main.py,sha256=VcYm3zXEAe_kT2DvReWmjtwaNMcQdF2Um24qXm6GEUo,24102
+aipr/providers.py,sha256=gyrGVFQYUhfoq3aTBnZEkcP9LY9e6mITZ39qt278yDk,6635
+aipr/prompts/__init__.py,sha256=VR8pQpDV0l-929daoPTAWa5zTdXfZ-_dmV2rZXhOu3s,140
+aipr/prompts/meta.xml,sha256=VLnvS85KEUG4Td0iqu4Bxs_qachMlHD1qqwFFuNx45M,2053
+aipr/prompts/prompts.py,sha256=SZYlSGDbJ__-_EJQTBLWJCHdDSj0ekws1ujR9trygXE,7693
+pr_generator_agent-1.0.0.dist-info/METADATA,sha256=tWo4MJYyPIAIr83nBQW_TDGPhg_53dc2tyOF1b5wKeE,4857
+pr_generator_agent-1.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pr_generator_agent-1.0.0.dist-info/entry_points.txt,sha256=Nx7jYZTtBlWrDKhHkPYD-UBs4xA-rr4Q-KTDzcpeIFU,40
+pr_generator_agent-1.0.0.dist-info/licenses/LICENSE,sha256=IeC3i1PYZv6yj4Pj6rp2Nzs2-F70buj7UTFsmLUJvO8,1070
+pr_generator_agent-1.0.0.dist-info/RECORD,,

pr_generator_agent-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pr_generator_agent-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+aipr = aipr.main:main

pr_generator_agent-1.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Daniel Scholl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.