llm.py

"""Provider-agnostic hosted-LLM client for the enhancer's optional "AI" engine.

Keys are supplied per-request by the user (never stored server-side). Supports
OpenAI-compatible providers (OpenAI, DeepSeek, OpenRouter, Ollama), Anthropic,
and Google Gemini. Every call is best-effort; failures raise so the caller can
fall back to the deterministic engine.
"""
import re

DEFAULT_MODELS = {
    "openai": "gpt-4o-mini",
    "anthropic": "claude-haiku-4-5-20251001",
    "deepseek": "deepseek-chat",
    "openrouter": "openai/gpt-4o-mini",
    "gemini": "gemini-1.5-flash",
    "ollama": "llama3.2",
}

# Providers that speak the OpenAI /chat/completions dialect.
_OPENAI_BASE = {
    "openai": "https://api.openai.com/v1",
    "deepseek": "https://api.deepseek.com/v1",
    "openrouter": "https://openrouter.ai/api/v1",
}

SUPPORTED = set(DEFAULT_MODELS)


def _system_prompt(language: str) -> str:
    return (
        f"You are a secure-coding assistant. Rewrite the following {language} code "
        "to fix every security vulnerability while preserving its behaviour. Return "
        "ONLY the corrected code, with no explanations, no commentary, and no "
        "markdown code fences."
    )


def strip_code_fences(text: str) -> str:
    text = (text or "").strip()
    fenced = re.match(r"^```[A-Za-z0-9_+-]*\n(.*)\n?```$", text, re.DOTALL)
    if fenced:
        return fenced.group(1).strip()
    text = re.sub(r"^```[A-Za-z0-9_+-]*\n?", "", text)
    text = re.sub(r"\n?```$", "", text)
    return text.strip()


def _ollama_base(base_url: str) -> str:
    base = (base_url or "http://localhost:11434/v1").strip().rstrip("/")
    if not re.match(r"^https?://", base):
        raise ValueError("Ollama base URL must start with http:// or https://")
    if not base.endswith("/v1"):
        base = base + "/v1"
    return base


def _openai_chat(base, api_key, model, system, user, max_tokens, timeout, referer=None):
    import requests

    headers = {"content-type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    if referer:  # OpenRouter likes these; harmless elsewhere.
        headers["HTTP-Referer"] = referer
        headers["X-Title"] = "FortiScan"
    resp = requests.post(
        f"{base}/chat/completions",
        headers=headers,
        json={
            "model": model,
            "max_tokens": max_tokens,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def _anthropic(api_key, model, system, user, max_tokens, timeout):
    import requests

    resp = requests.post(
        "https://api.anthropic.com/v1/messages",
        headers={
            "x-api-key": api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        },
        json={
            "model": model,
            "max_tokens": max_tokens,
            "system": system,
            "messages": [{"role": "user", "content": user}],
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    return resp.json()["content"][0]["text"]


def _gemini(api_key, model, system, user, max_tokens, timeout):
    import requests

    url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
    resp = requests.post(
        url,
        headers={"content-type": "application/json", "x-goog-api-key": api_key},
        json={
            "system_instruction": {"parts": [{"text": system}]},
            "contents": [{"parts": [{"text": user}]}],
            "generationConfig": {"maxOutputTokens": max_tokens},
        },
        timeout=timeout,
    )
    resp.raise_for_status()
    data = resp.json()
    return data["candidates"][0]["content"]["parts"][0]["text"]


def call_llm(provider, api_key="", model="", code="", language="python",
             base_url="", max_tokens=2048, timeout=45):
    """Return the LLM's rewritten code (fences stripped). Raises on failure."""
    provider = (provider or "").strip().lower()
    if provider not in SUPPORTED:
        raise ValueError(f"Unsupported provider: {provider!r}")
    model = (model or "").strip() or DEFAULT_MODELS[provider]
    system = _system_prompt(language)

    if provider in _OPENAI_BASE:
        referer = "https://fortiscan.app" if provider == "openrouter" else None
        text = _openai_chat(_OPENAI_BASE[provider], api_key, model, system, code,
                            max_tokens, timeout, referer=referer)
    elif provider == "ollama":
        text = _openai_chat(_ollama_base(base_url), "", model, system, code, max_tokens, timeout)
    elif provider == "anthropic":
        text = _anthropic(api_key, model, system, code, max_tokens, timeout)
    elif provider == "gemini":
        text = _gemini(api_key, model, system, code, max_tokens, timeout)
    else:  # pragma: no cover - guarded above
        raise ValueError(f"Unsupported provider: {provider!r}")

    return strip_code_fences(text)


def test_llm(provider, api_key="", model="", base_url="", timeout=15):
    """Live connectivity/auth check. Returns (ok: bool, message: str)."""
    import requests

    provider = (provider or "").strip().lower()
    if provider not in SUPPORTED:
        return False, f"Unsupported provider: {provider}"
    if provider != "ollama" and not api_key:
        return False, "An API key is required for this provider."

    model = (model or "").strip() or DEFAULT_MODELS[provider]
    try:
        # Minimal generation verifies connectivity, auth, and model availability.
        call_llm(provider, api_key=api_key, model=model, code="print(1)",
                 language="python", base_url=base_url, max_tokens=16, timeout=timeout)
        return True, f"Connected — {provider} / {model} is live."
    except requests.HTTPError as e:
        status = e.response.status_code if e.response is not None else "?"
        detail = ""
        if e.response is not None:
            try:
                body = e.response.json()
                detail = (body.get("error", {}) or {}).get("message", "") if isinstance(body.get("error"), dict) else str(body.get("error", ""))
            except Exception:
                detail = e.response.text[:200]
        hint = " (check the API key)" if status in (401, 403) else " (check the model name)" if status == 404 else ""
        return False, f"HTTP {status}{hint}: {detail or 'request rejected'}"
    except requests.ConnectionError:
        return False, "Connection failed — is the endpoint reachable? (For Ollama, is it running?)"
    except requests.Timeout:
        return False, "Timed out contacting the provider."
    except Exception as e:
        return False, f"{e.__class__.__name__}: {e}"