switchforge 1.0.0__py3-none-any.whl

forge_core/__init__.py

@@ -0,0 +1,3 @@
+"""Forge Core — AI-powered backend test generation engine."""
+
+__version__ = "2.0.0"

forge_core/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for `python -m forge_core`."""
+
+from forge_core.cli import app
+
+if __name__ == "__main__":
+    app()

forge_core/ai/prompts.py

@@ -0,0 +1,87 @@
+"""Load and template .md prompt files for AI consumption."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from forge_core.utils import logger
+
+
+def load_prompt(prompts_dir: Path, prompt_name: str) -> str:
+    """Load a prompt .md file from the prompts directory.
+
+    Args:
+        prompts_dir: Path to the prompts directory.
+        prompt_name: Name of the prompt file (e.g., 'detect-tech-stack').
+
+    Returns:
+        The prompt content as a string.
+    """
+    file_path = prompts_dir / f"{prompt_name}.prompt.md"
+    if not file_path.exists():
+        logger.warn(f"Prompt not found: {file_path}")
+        return ""
+
+    content = file_path.read_text(encoding="utf-8")
+    # Strip YAML frontmatter if present
+    if content.startswith("---"):
+        parts = content.split("---", 2)
+        if len(parts) >= 3:
+            content = parts[2].strip()
+
+    return content
+
+
+def template_prompt(prompt: str, variables: dict[str, str]) -> str:
+    """Replace template variables in a prompt.
+
+    Variables use {{variable_name}} syntax.
+    """
+    result = prompt
+    for key, value in variables.items():
+        result = result.replace(f"{{{{{key}}}}}", value)
+    return result
+
+
+def build_file_context(files: dict[str, str], max_files: int = 50) -> str:
+    """Build a file context string from a dict of {path: content}.
+
+    Formats each file with clear separators for the AI to parse.
+    """
+    parts: list[str] = []
+    for i, (path, content) in enumerate(files.items()):
+        if i >= max_files:
+            parts.append(f"\n... and {len(files) - max_files} more files (truncated)")
+            break
+        parts.append(f"--- {path} ---\n{content}")
+
+    return "\n\n".join(parts)
+
+
+def load_learnings(central_path: str | None, local_path: Path | None) -> str:
+    """Load LEARNINGS.md from central hub and/or local project."""
+    learnings_parts: list[str] = []
+
+    if central_path:
+        central_file = Path(central_path) / "LEARNINGS.md"
+        if central_file.exists():
+            learnings_parts.append(
+                f"=== Central Learnings ===\n{central_file.read_text(encoding='utf-8')}"
+            )
+
+    if local_path:
+        local_file = local_path / "LEARNINGS.md"
+        if local_file.exists():
+            learnings_parts.append(
+                f"=== Project Learnings ===\n{local_file.read_text(encoding='utf-8')}"
+            )
+
+    return "\n\n".join(learnings_parts)
+
+
+def load_knowledge_pack(central_path: str, pack_name: str) -> str:
+    """Load a specific knowledge pack (e.g., kotlin-ktor.md)."""
+    pack_path = Path(central_path) / "knowledge-packs" / f"{pack_name}.md"
+    if pack_path.exists():
+        return pack_path.read_text(encoding="utf-8")
+    return ""

forge_core/ai/provider.py

@@ -0,0 +1,121 @@
+"""LiteLLM-based AI provider — unified interface for 100+ models."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import litellm
+
+from forge_core.models.config import AIConfig, AIProvider
+from forge_core.utils import logger
+from forge_core.utils.tokens import count_tokens
+
+# Suppress LiteLLM's verbose logging
+litellm.suppress_debug_info = True
+
+
+def _resolve_model(config: AIConfig) -> str:
+    """Resolve the model string for LiteLLM based on provider."""
+    if config.provider == AIProvider.OLLAMA:
+        return f"ollama/{config.model}"
+    if config.provider == AIProvider.AZURE:
+        return f"azure/{config.model}"
+    if config.provider == AIProvider.BEDROCK:
+        return f"bedrock/{config.model}"
+    return config.model
+
+
+def complete(
+    config: AIConfig,
+    system_prompt: str,
+    user_prompt: str,
+    json_mode: bool = False,
+    max_tokens: int | None = None,
+) -> str:
+    """Send a completion request to the configured AI provider.
+
+    Args:
+        config: AI provider configuration.
+        system_prompt: System-level instructions.
+        user_prompt: User message with code/context.
+        json_mode: If True, request JSON response format.
+        max_tokens: Override max tokens for this call.
+
+    Returns:
+        The AI's response text.
+    """
+    model = _resolve_model(config)
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+
+    kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+        "temperature": config.temperature,
+        "max_tokens": max_tokens or config.max_tokens,
+    }
+
+    if config.api_key:
+        kwargs["api_key"] = config.api_key
+    if config.base_url:
+        kwargs["api_base"] = config.base_url
+    if json_mode:
+        kwargs["response_format"] = {"type": "json_object"}
+
+    # Log token usage
+    input_tokens = count_tokens(system_prompt + user_prompt, config.model)
+    logger.info(f"AI call → {model} ({input_tokens} input tokens)")
+
+    try:
+        response = litellm.completion(**kwargs)
+        content = response.choices[0].message.content or ""
+
+        output_tokens = count_tokens(content, config.model)
+        logger.info(f"AI response ← {output_tokens} output tokens")
+
+        return content
+    except Exception as e:
+        logger.error(f"AI call failed: {e}")
+        raise
+
+
+def complete_with_fallback(
+    config: AIConfig,
+    system_prompt: str,
+    user_prompt: str,
+    fallback_models: list[str] | None = None,
+    json_mode: bool = False,
+) -> str:
+    """Try primary model, fall back to alternatives on failure.
+
+    Useful for complex classes where one model might fail.
+    """
+    models = [_resolve_model(config)] + (fallback_models or [])
+
+    last_error = None
+    for model in models:
+        try:
+            kwargs: dict[str, Any] = {
+                "model": model,
+                "messages": [
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt},
+                ],
+                "temperature": config.temperature,
+                "max_tokens": config.max_tokens,
+            }
+            if config.api_key:
+                kwargs["api_key"] = config.api_key
+            if json_mode:
+                kwargs["response_format"] = {"type": "json_object"}
+
+            response = litellm.completion(**kwargs)
+            return response.choices[0].message.content or ""
+        except Exception as e:
+            last_error = e
+            logger.warn(f"Model {model} failed, trying next: {e}")
+            continue
+
+    raise RuntimeError(f"All models failed. Last error: {last_error}")

forge_core/ai/structured.py

@@ -0,0 +1,68 @@
+"""Structured AI outputs using instructor + pydantic."""
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+import instructor
+import litellm
+from pydantic import BaseModel
+
+from forge_core.models.config import AIConfig
+from forge_core.ai.provider import _resolve_model
+from forge_core.utils import logger
+
+T = TypeVar("T", bound=BaseModel)
+
+
+def extract(
+    config: AIConfig,
+    system_prompt: str,
+    user_prompt: str,
+    response_model: type[T],
+    max_retries: int = 2,
+) -> T:
+    """Extract structured data from AI using instructor.
+
+    Forces the AI to return data matching a pydantic model schema.
+    Retries on validation failure.
+
+    Args:
+        config: AI provider configuration.
+        system_prompt: System-level instructions.
+        user_prompt: User message with code/context.
+        response_model: Pydantic model class to extract.
+        max_retries: Number of retries on validation failure.
+
+    Returns:
+        An instance of response_model populated by the AI.
+    """
+    model = _resolve_model(config)
+
+    client = instructor.from_litellm(litellm.completion)
+
+    kwargs: dict[str, Any] = {
+        "model": model,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        "temperature": config.temperature,
+        "max_retries": max_retries,
+        "response_model": response_model,
+    }
+
+    if config.api_key:
+        kwargs["api_key"] = config.api_key
+    if config.base_url:
+        kwargs["api_base"] = config.base_url
+
+    logger.info(f"Structured extraction → {model} → {response_model.__name__}")
+
+    try:
+        result = client.chat.completions.create(**kwargs)
+        logger.success(f"Extracted {response_model.__name__} successfully")
+        return result
+    except Exception as e:
+        logger.error(f"Structured extraction failed: {e}")
+        raise

forge_core/auth.py

@@ -0,0 +1,72 @@
+"""SaaS API authentication and license verification."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+from forge_core.models.config import ForgeConfig, Plan, PlanLimits
+from forge_core.utils import logger
+from forge_core.utils.reporter import check_license
+
+
+def verify_license(config: ForgeConfig) -> ForgeConfig:
+    """Verify license with SaaS API and update plan limits.
+
+    If no auth token is set, runs in free/offline mode.
+    """
+    if not config.auth_token:
+        logger.info("Running in offline mode (no auth token)")
+        config.limits = PlanLimits.for_plan(Plan.FREE)
+        return config
+
+    # Check with SaaS API
+    logger.info("Verifying license with SaaS API...")
+    try:
+        license_info = asyncio.run(check_license(config))
+    except Exception:
+        license_info = {}
+
+    if not license_info:
+        logger.warn("Could not verify license — falling back to free tier limits")
+        config.limits = PlanLimits.for_plan(Plan.FREE)
+        return config
+
+    # Update config from license response
+    plan_str = license_info.get("plan", "free")
+    try:
+        plan = Plan(plan_str)
+    except ValueError:
+        plan = Plan.FREE
+
+    config.limits = PlanLimits.for_plan(plan)
+    config.tenant.org_id = license_info.get("org_id", config.tenant.org_id)
+    config.tenant.org_name = license_info.get("org_name", config.tenant.org_name)
+
+    # If Pro/Enterprise, use SaaS proxy for AI calls
+    if plan in (Plan.PRO, Plan.ENTERPRISE):
+        config.ai.use_saas_proxy = True
+        if not config.ai.api_key:
+            config.ai.api_key = config.auth_token  # SaaS proxy uses auth token
+            config.ai.base_url = f"{config.saas_api_url}/api/v1/ai"
+
+    logger.success(f"License verified: {plan.value} plan ({config.tenant.org_name})")
+    return config
+
+
+def save_auth_token(token: str) -> None:
+    """Save auth token to ~/.forge-core/auth."""
+    auth_dir = Path.home() / ".forge-core"
+    auth_dir.mkdir(exist_ok=True)
+    auth_file = auth_dir / "auth"
+    auth_file.write_text(token, encoding="utf-8")
+    auth_file.chmod(0o600)
+    logger.success("Auth token saved")
+
+
+def load_auth_token() -> str:
+    """Load auth token from ~/.forge-core/auth."""
+    auth_file = Path.home() / ".forge-core" / "auth"
+    if auth_file.exists():
+        return auth_file.read_text(encoding="utf-8").strip()
+    return ""

forge_core/cli.py

@@ -0,0 +1,209 @@
+"""Forge Core CLI — powered by typer."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich.console import Console
+
+from forge_core import __version__
+
+app = typer.Typer(
+    name="forge-core",
+    help="AI-powered backend test generation engine",
+    no_args_is_help=True,
+)
+console = Console()
+
+
+@app.command()
+def run(
+    project_path: Path = typer.Argument(
+        ".",
+        help="Path to the backend project to test",
+        exists=True,
+        dir_okay=True,
+        file_okay=False,
+    ),
+    target_coverage: float = typer.Option(
+        90.0, "--target", "-t", help="Target line coverage percentage"
+    ),
+    max_iterations: int = typer.Option(
+        10, "--max-iterations", "-i", help="Maximum generation iterations"
+    ),
+    model: str = typer.Option(
+        "", "--model", "-m", help="AI model to use (e.g., gpt-4o, claude-sonnet-4-20250514)"
+    ),
+    provider: str = typer.Option(
+        "", "--provider", "-p", help="AI provider (openai, anthropic, ollama, azure)"
+    ),
+    api_key: str = typer.Option(
+        "", "--api-key", "-k", help="AI provider API key (or set OPENAI_API_KEY env var)"
+    ),
+    mode: str = typer.Option(
+        "full", "--mode", help="Run mode: full, targeted, analyze_only, analyze_review"
+    ),
+    files: Optional[list[str]] = typer.Option(
+        None, "--file", "-f", help="Specific files to test (targeted mode)"
+    ),
+):
+    """Run Forge Core on a backend project to generate unit tests."""
+    from forge_core.auth import load_auth_token, verify_license
+    from forge_core.config import load_config
+    from forge_core.models.config import RunMode
+    from forge_core.orchestrator import Orchestrator
+    from forge_core.utils import logger
+    from forge_core.utils.reporter import upload_report
+
+    # Resolve project path
+    project_path = project_path.resolve()
+
+    # Load config
+    auth_token = load_auth_token()
+    config = load_config(
+        project_path=project_path,
+        api_key=api_key,
+        provider=provider,
+        model=model,
+        target_coverage=target_coverage,
+        auth_token=auth_token,
+    )
+    config.max_iterations = max_iterations
+
+    # Set mode
+    try:
+        config.mode = RunMode(mode)
+    except ValueError:
+        logger.error(f"Invalid mode: {mode}. Use: full, targeted, analyze_only, analyze_review")
+        raise typer.Exit(1)
+
+    if files:
+        config.mode = RunMode.TARGETED
+        config.target_files = list(files)
+
+    # Verify license
+    config = verify_license(config)
+
+    # Validate API key
+    if not config.ai.api_key:
+        console.print(
+            "[red]No API key found.[/red]\n\n"
+            "Set one of:\n"
+            "  export OPENAI_API_KEY=sk-...\n"
+            "  export ANTHROPIC_API_KEY=sk-ant-...\n"
+            "  forge-core run --api-key sk-...\n"
+            "  forge-core login (for Pro/Enterprise)\n"
+        )
+        raise typer.Exit(1)
+
+    # Run the pipeline
+    orchestrator = Orchestrator(config)
+    report = orchestrator.run()
+
+    # Upload report to SaaS (if authenticated)
+    if config.auth_token:
+        asyncio.run(upload_report(config, report.model_dump()))
+
+    # Exit with appropriate code
+    if report.production_files_changed > 0:
+        logger.error("SAFETY VIOLATION: Production files were changed!")
+        raise typer.Exit(2)
+
+    raise typer.Exit(0)
+
+
+@app.command()
+def login(
+    token: str = typer.Option(
+        "", "--token", "-t", help="Auth token from theswitchcompany.online"
+    ),
+):
+    """Authenticate with TheSwitchCompany SaaS for Pro/Enterprise features."""
+    from forge_core.auth import save_auth_token, verify_license
+    from forge_core.config import load_config
+    from forge_core.utils import logger
+
+    if not token:
+        console.print(
+            "Get your auth token from:\n"
+            "  [blue]https://theswitchcompany.online/dashboard/settings[/blue]\n\n"
+            "Then run:\n"
+            "  forge-core login --token YOUR_TOKEN"
+        )
+        raise typer.Exit(0)
+
+    save_auth_token(token)
+
+    # Verify the token
+    config = load_config(Path("."), auth_token=token)
+    config = verify_license(config)
+    logger.success(f"Logged in as {config.tenant.org_name or 'user'}")
+
+
+@app.command()
+def init(
+    project_path: Path = typer.Argument(".", help="Path to the project"),
+):
+    """Initialize Forge Core in a project (creates .github/agent-config.yml)."""
+    from forge_core.utils import logger
+
+    project_path = project_path.resolve()
+    config_dir = project_path / ".github"
+    config_dir.mkdir(exist_ok=True)
+
+    config_file = config_dir / "agent-config.yml"
+    if config_file.exists():
+        logger.info("agent-config.yml already exists")
+        raise typer.Exit(0)
+
+    config_file.write_text(
+        "# Forge Core configuration\n"
+        "# See: https://theswitchcompany.online/docs/forge-core\n\n"
+        "# Tenant (populated by SaaS or set manually)\n"
+        'org_id: ""\n'
+        'user_id: ""\n'
+        'project_id: ""\n\n'
+        "# Runtime\n"
+        'runtime: "auto"\n'
+        "max_parallel_agents: 4\n"
+        'cache_dir: ".forge-cache"\n',
+        encoding="utf-8",
+    )
+
+    logger.success(f"Initialized Forge Core in {project_path}")
+    console.print("\nNext: [blue]forge-core run[/blue] to generate tests")
+
+
+@app.command()
+def version():
+    """Show Forge Core version."""
+    console.print(f"Forge Core v{__version__}")
+
+
+@app.command()
+def status():
+    """Show current configuration and license status."""
+    from forge_core.auth import load_auth_token, verify_license
+    from forge_core.config import load_config
+
+    auth_token = load_auth_token()
+    config = load_config(Path("."), auth_token=auth_token)
+
+    if auth_token:
+        config = verify_license(config)
+
+    console.print(f"\n[bold]Forge Core v{__version__}[/bold]\n")
+    console.print(f"  Plan:     {config.limits.plan.value}")
+    console.print(f"  Org:      {config.tenant.org_name or '(not set)'}")
+    console.print(f"  Model:    {config.ai.model}")
+    console.print(f"  Provider: {config.ai.provider.value}")
+    console.print(f"  API Key:  {'***' + config.ai.api_key[-4:] if config.ai.api_key else '(not set)'}")
+    console.print(f"  Auth:     {'✓ logged in' if auth_token else '✗ offline'}")
+    console.print()
+
+
+if __name__ == "__main__":
+    app()

forge_core/config.py

@@ -0,0 +1,118 @@
+"""Configuration loader — merges agent-config.yml, env vars, and CLI args."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import yaml
+
+from forge_core.models.config import AIConfig, AIProvider, ForgeConfig, TenantInfo
+from forge_core.utils import logger
+
+
+def load_config(
+    project_path: Path,
+    api_key: str = "",
+    provider: str = "",
+    model: str = "",
+    target_coverage: float = 0,
+    auth_token: str = "",
+) -> ForgeConfig:
+    """Load configuration from agent-config.yml, env vars, and CLI overrides.
+
+    Priority: CLI args > env vars > agent-config.yml > defaults.
+    """
+    config = ForgeConfig(project_path=project_path)
+
+    # 1. Load agent-config.yml from the project
+    yml_path = project_path / ".github" / "agent-config.yml"
+    if yml_path.exists():
+        _load_yml(config, yml_path)
+
+    # 2. Environment variables override
+    _load_env(config)
+
+    # 3. CLI args override
+    if api_key:
+        config.ai.api_key = api_key
+    if provider:
+        config.ai.provider = AIProvider(provider)
+    if model:
+        config.ai.model = model
+    if target_coverage > 0:
+        config.target_coverage = target_coverage
+    if auth_token:
+        config.auth_token = auth_token
+
+    # 4. Resolve prompts directory
+    config.prompts_dir = str(project_path / ".github" / "prompts")
+
+    return config
+
+
+def _load_yml(config: ForgeConfig, yml_path: Path) -> None:
+    """Load settings from agent-config.yml."""
+    try:
+        data = yaml.safe_load(yml_path.read_text(encoding="utf-8")) or {}
+    except Exception as e:
+        logger.warn(f"Failed to parse {yml_path}: {e}")
+        return
+
+    if "central_agent_path" in data:
+        config.central_agent_path = data["central_agent_path"]
+    if "knowledge_packs_dir" in data:
+        config.knowledge_packs_dir = data["knowledge_packs_dir"]
+    if "cache_dir" in data:
+        config.cache_dir = data["cache_dir"]
+
+    # Tenant info
+    config.tenant = TenantInfo(
+        org_id=data.get("org_id", ""),
+        org_name=data.get("org_name", ""),
+        user_id=data.get("user_id", ""),
+        project_id=data.get("project_id", ""),
+    )
+
+    # Runtime config
+    if "runtime" in data and data["runtime"] != "auto":
+        try:
+            config.ai.provider = AIProvider(data["runtime"])
+        except ValueError:
+            pass
+    if "max_parallel_agents" in data:
+        config.limits.max_parallel_agents = int(data["max_parallel_agents"])
+
+
+def _load_env(config: ForgeConfig) -> None:
+    """Load settings from environment variables."""
+    env_map = {
+        "FORGE_API_KEY": "api_key",
+        "OPENAI_API_KEY": "api_key",
+        "ANTHROPIC_API_KEY": "api_key",
+        "FORGE_MODEL": "model",
+        "FORGE_PROVIDER": "provider",
+        "FORGE_AUTH_TOKEN": "auth_token",
+    }
+
+    for env_var, field in env_map.items():
+        value = os.environ.get(env_var, "")
+        if not value:
+            continue
+
+        if field == "api_key" and not config.ai.api_key:
+            config.ai.api_key = value
+            # Auto-detect provider from key prefix
+            if env_var == "ANTHROPIC_API_KEY" or value.startswith("sk-ant-"):
+                config.ai.provider = AIProvider.ANTHROPIC
+                if config.ai.model == "gpt-4o":
+                    config.ai.model = "claude-sonnet-4-20250514"
+        elif field == "model":
+            config.ai.model = value
+        elif field == "provider":
+            try:
+                config.ai.provider = AIProvider(value)
+            except ValueError:
+                pass
+        elif field == "auth_token":
+            config.auth_token = value