rigdoctor/src/rigdoctor/core/ai.py

"""AI assistant (M14, D24): explain the collected diagnostics in plain language.

**Strictly opt-in and never automatic** — the model is contacted ONLY from a direct user
action ("Explain with AI" / ``rigdoctor ai explain``), never on launch, after a diagnostic, or
in any loop. Choosing/configuring a provider does not contact anything. The user must pick a
provider explicitly (there is no default).

Two providers, both over stdlib ``urllib`` (no pip deps in core):
  * **ollama** — a local server (data stays on the machine, no key).
  * **claude** — the Anthropic Messages API (key in the keyring).

Answers are *grounded*: we pass the actual findings plus matched reference facts
(:mod:`ai_knowledge`) and ask the model to reason over them. Output is advisory (D9).
"""

from __future__ import annotations

import json
import re
import urllib.error
import urllib.request

from .. import config
from . import ai_knowledge

_APPID_RE = re.compile(r"\b\d{5,7}\b")  # Steam app IDs are 5–7 digits

PROVIDERS = ("ollama", "claude")
OLLAMA_DEFAULT_ENDPOINT = "http://localhost:11434"
# Suggested Ollama model — strong instruction-following that fits an 8 GB GPU at Q4. Because we
# ground the prompt with reference facts, a 7B model is sufficient here.
OLLAMA_SUGGESTED_MODEL = "qwen2.5:7b"
CLAUDE_ENDPOINT = "https://api.anthropic.com/v1/messages"
CLAUDE_DEFAULT_MODEL = "claude-opus-4-7"
CLAUDE_MAX_TOKENS = 2000
ANTHROPIC_VERSION = "2023-06-01"

SYSTEM_PROMPT = (
    "You are RigDoctor's hardware-diagnostics assistant for Linux gamers (Ubuntu + NVIDIA, games "
    "via Steam/Proton). You are given session context, the structured findings RigDoctor "
    "collected — which may include recent game/Proton/system log excerpts scoped to this session "
    "— plus reference facts. Use the GAME NAME from the session context; never guess the game "
    "from log paths or app IDs. Correlate log errors with the findings to pinpoint WHEN and WHY "
    "things went wrong, identify the most likely root cause, and give concrete, ordered next "
    "steps with exact Linux commands where useful.\n"
    "Rules: Base your reasoning ONLY on the data and reference facts provided — never invent "
    "readings, hardware, or log lines. This is LINUX: never suggest Windows-only steps (e.g. "
    "'run as administrator', registry edits, toggling antivirus). Treat log lines flagged BENIGN "
    "in the reference facts as non-causal. If no crash was recorded and there are no warning or "
    "critical findings, say plainly that the session looks healthy and do NOT manufacture a "
    "problem. Be concise. Present fixes as suggestions and warn before anything that risks data "
    "loss or instability. Format your answer in Markdown."
)


def provider() -> str:
    return config.load_config().get("ai_provider", "")


def model() -> str:
    m = config.load_config().get("ai_model", "").strip()
    if m:
        return m
    return CLAUDE_DEFAULT_MODEL if provider() == "claude" else ""


def endpoint() -> str:
    ep = config.load_config().get("ai_endpoint", OLLAMA_DEFAULT_ENDPOINT).strip()
    return ep or OLLAMA_DEFAULT_ENDPOINT


def is_local() -> bool:
    return provider() == "ollama"


def is_configured() -> bool:
    """Whether the chosen provider is ready (does NOT contact anything)."""
    p = provider()
    if p == "claude":
        return bool(config.load_ai_key())
    if p == "ollama":
        return bool(model())  # a model name is required; endpoint has a default
    return False  # no provider chosen


def provider_label() -> str:
    p = provider()
    if p == "claude":
        return f"Claude ({model()})"
    if p == "ollama":
        return f"Ollama ({model() or '?'} @ {endpoint()})"
    return "not configured"


def appid_glossary(text: str) -> str:
    """Resolve Steam app IDs that appear in `text` against the user's scanned library.

    We don't teach the model app IDs — we look them up locally and hand it the mapping, so it
    names games correctly instead of guessing. Only IDs we can resolve are listed.
    """
    candidates = set(_APPID_RE.findall(text))
    if not candidates:
        return ""
    try:
        from . import steam
        names = steam.appid_names()
    except Exception:  # never let a glossary lookup break an explanation
        return ""
    known = sorted((i, names[i]) for i in candidates if i in names)
    if not known:
        return ""
    return "App IDs (resolved from your installed games):\n" + "\n".join(
        f"- {appid} = {name}" for appid, name in known)


def build_prompt(findings_text: str) -> str:
    """The user-message content: app-ID glossary + matched reference facts + the findings."""
    parts = []
    glossary = appid_glossary(findings_text)
    if glossary:
        parts.append(glossary)
        parts.append("")
    facts = ai_knowledge.relevant(findings_text)
    if facts:
        parts.append("Reference facts (use these to interpret the findings):")
        parts += [f"- {f}" for f in facts]
        parts.append("")
    parts.append("Collected findings:")
    parts.append(findings_text.strip() or "(no findings provided)")
    return "\n".join(parts)


def explain(findings_text: str, timeout: float = 120.0) -> tuple[bool, str]:
    """Contact the configured provider to explain the findings. Returns (ok, text | error).

    The caller MUST be a direct user action (D24) — this never runs automatically.
    """
    content = build_prompt(findings_text)
    try:
        if provider() == "claude":
            return _claude(content, timeout)
        if provider() == "ollama":
            return _ollama(content, timeout)
        return False, "No AI provider is configured (Settings → AI assistant)."
    except urllib.error.HTTPError as exc:
        return False, _http_error(exc)
    except (urllib.error.URLError, OSError, TimeoutError) as exc:
        return False, f"Couldn't reach the AI provider: {exc}"
    except (ValueError, KeyError, IndexError) as exc:
        return False, f"Unexpected response from the AI provider: {exc}"


def explain_stream(findings_text: str, on_chunk, timeout: float = 180.0) -> tuple[bool, str]:
    """Like :func:`explain`, but calls ``on_chunk(text_delta)`` as tokens arrive and returns
    ``(ok, full_text)`` at the end. Caller MUST be a direct user action (D24)."""
    content = build_prompt(findings_text)
    try:
        if provider() == "claude":
            return _claude_stream(content, on_chunk, timeout)
        if provider() == "ollama":
            return _ollama_stream(content, on_chunk, timeout)
        return False, "No AI provider is configured (Settings → AI assistant)."
    except urllib.error.HTTPError as exc:
        return False, _http_error(exc)
    except (urllib.error.URLError, OSError, TimeoutError) as exc:
        return False, f"Couldn't reach the AI provider: {exc}"
    except (ValueError, KeyError, IndexError) as exc:
        return False, f"Unexpected response from the AI provider: {exc}"


def _post(url: str, payload: dict, headers: dict, timeout: float) -> dict:
    req = urllib.request.Request(
        url, data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", **headers},
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.load(resp)


def _ollama(content: str, timeout: float) -> tuple[bool, str]:
    if not model():
        return False, "No Ollama model is set (Settings → AI assistant)."
    payload = {"model": model(), "system": SYSTEM_PROMPT, "prompt": content, "stream": False}
    out = _post(endpoint().rstrip("/") + "/api/generate", payload, {}, timeout)
    return True, (out.get("response") or "").strip() or "(the model returned an empty response)"


def _claude(content: str, timeout: float) -> tuple[bool, str]:
    key = config.load_ai_key()
    if not key:
        return False, "No Claude API key is set (Settings → AI assistant)."
    # One-shot call: no prompt caching (single request, short system prompt) and no thinking
    # (keeps a button-press snappy). Sampling params are omitted (removed on current Opus).
    payload = {
        "model": model(),
        "max_tokens": CLAUDE_MAX_TOKENS,
        "system": SYSTEM_PROMPT,
        "messages": [{"role": "user", "content": content}],
    }
    headers = {"x-api-key": key, "anthropic-version": ANTHROPIC_VERSION}
    out = _post(CLAUDE_ENDPOINT, payload, headers, timeout)
    text = "\n".join(b.get("text", "") for b in out.get("content", []) if b.get("type") == "text")
    return True, text.strip() or "(the model returned no text)"


def _stream_request(url: str, payload: dict, headers: dict, timeout: float):
    req = urllib.request.Request(
        url, data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", **headers})
    return urllib.request.urlopen(req, timeout=timeout)


def _ollama_stream(content: str, on_chunk, timeout: float) -> tuple[bool, str]:
    if not model():
        return False, "No Ollama model is set (Settings → AI assistant)."
    payload = {"model": model(), "system": SYSTEM_PROMPT, "prompt": content, "stream": True}
    parts: list[str] = []
    with _stream_request(endpoint().rstrip("/") + "/api/generate", payload, {}, timeout) as resp:
        for raw in resp:  # newline-delimited JSON objects
            line = raw.decode("utf-8", "replace").strip()
            if not line:
                continue
            obj = json.loads(line)
            chunk = obj.get("response", "")
            if chunk:
                parts.append(chunk)
                on_chunk(chunk)
            if obj.get("done"):
                break
    return True, "".join(parts).strip() or "(the model returned an empty response)"


def _claude_stream(content: str, on_chunk, timeout: float) -> tuple[bool, str]:
    key = config.load_ai_key()
    if not key:
        return False, "No Claude API key is set (Settings → AI assistant)."
    payload = {
        "model": model(), "max_tokens": CLAUDE_MAX_TOKENS, "system": SYSTEM_PROMPT,
        "messages": [{"role": "user", "content": content}], "stream": True,
    }
    headers = {"x-api-key": key, "anthropic-version": ANTHROPIC_VERSION}
    parts: list[str] = []
    with _stream_request(CLAUDE_ENDPOINT, payload, headers, timeout) as resp:
        for raw in resp:  # SSE: parse `data:` lines, accumulate text deltas
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data:"):
                continue
            try:
                event = json.loads(line[5:].strip())
            except ValueError:
                continue
            etype = event.get("type")
            if etype == "content_block_delta" and event.get("delta", {}).get("type") == "text_delta":
                chunk = event["delta"].get("text", "")
                if chunk:
                    parts.append(chunk)
                    on_chunk(chunk)
            elif etype == "error":
                return False, event.get("error", {}).get("message", "stream error")
            elif etype == "message_stop":
                break
    return True, "".join(parts).strip() or "(the model returned no text)"


def _http_error(exc: urllib.error.HTTPError) -> str:
    detail = ""
    try:
        body = exc.read().decode("utf-8", "replace")
        detail = json.loads(body).get("error", {}).get("message", "") or ""
    except (ValueError, OSError):
        pass
    hint = " — check your API key in Settings → AI assistant." if exc.code in (401, 403) else ""
    return f"AI request failed (HTTP {exc.code}){hint}{(': ' + detail) if detail else ''}"


def format_findings(findings, header: str = "") -> str:
    """Render M4 Finding objects (or similar) into the plain-text block we send the model."""
    lines = [header] if header else []
    for f in findings:
        severity = str(getattr(f, "severity", "")).upper()
        category = getattr(f, "category", "")
        title = getattr(f, "title", "")
        detail = getattr(f, "detail", "")
        line = f"- [{severity}] {category}: {title}".rstrip()
        if detail:
            line += f" — {detail}"
        lines.append(line)
    return "\n".join(lines) if lines else "No findings."