src/llm.py

"""
llm.py — the LLM slot (Task 2 wiring point).

This is where the Qwen2.5-3B-Instruct GGUF model lives. The module is written
now so the rest of the pipeline has a clean, lazy single point of integration.

Until the model is wired in, `is_available()` returns False and the pipeline
falls back to the heuristic parser in `spec.interpret_query`.
"""

from __future__ import annotations

import os
import logging
from typing import Optional

from .spec import QuantumSpec, interpret_query, parse_dsl

log = logging.getLogger(__name__)

# ── Configuration ────────────────────────────────────────────────────────────
MODEL_REPO = os.environ.get(
    "QAC_L_MODEL_REPO",
    "YOUR_USERNAME/qwen3b-quantum-compiler", # Replace with your repo
)
MODEL_FILE = os.environ.get(
    "QAC_L_MODEL_FILE",
    "qwen3b_quantum.gguf",
)
N_CTX = int(os.environ.get("QAC_L_N_CTX", "4096"))
N_THREADS = int(os.environ.get("QAC_L_N_THREADS", "2"))

_LLM = None
_LOAD_ATTEMPTED = False

# ── Loading ──────────────────────────────────────────────────────────────────

def _load():
    """Lazy-load the GGUF model. Returns the llama-cpp instance or None."""
    global _LLM, _LOAD_ATTEMPTED
    if _LOAD_ATTEMPTED:
        return _LLM
    _LOAD_ATTEMPTED = True
    try:
        from huggingface_hub import hf_hub_download
        from llama_cpp import Llama

        path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        log.info("Loading LLM %s/%s (ctx=%d, threads=%d)",
                 MODEL_REPO, MODEL_FILE, N_CTX, N_THREADS)
        _LLM = Llama(
            model_path=path,
            n_ctx=N_CTX,
            n_threads=N_THREADS,
            verbose=False,
        )
        log.info("LLM loaded.")
    except Exception as exc:
        log.warning("LLM unavailable, falling back to heuristic parser: %s", exc)
        _LLM = None
    return _LLM

def is_available() -> bool:
    """True iff the GGUF model can be loaded in this environment."""
    return _load() is not None

# ── Generation ───────────────────────────────────────────────────────────────

_COMPILE_SYS = (
    "You are a quantum compiler. Translate the user's natural-language request "
    "into exactly one strict DSL line. Valid formats:\n"
    "  [ACTION: RANDOM] [QUBITS: n]          # n in 1..5\n"
    "  [ACTION: VQE]    [DISTANCE: r]        # r in 0.5..2.5 (Å)\n"
    "  [ACTION: BELL]   [PAIRS: k]           # k in 1..4\n"
    "  [ACTION: GROVER] [ITEMS: m]           # m in 2..16\n"
    "  [ACTION: QAOA]   [NODES: v] [DEPTH: p]# v in 2..8, p in 1..3\n"
    "Output ONLY the DSL line, nothing else."
)

_REPORT_SYS = (
    "You are a senior quantum physicist. Given raw measurement counts and "
    "computed statistics, explain the result in clear, scientifically accurate "
    "English. Treat probabilities below 5% as expected quantum noise. Be concise."
)

def _chat(llm, system: str, user: str, *, max_tokens: int, temperature: float) -> str:
    resp = llm.create_chat_completion(
        messages=[{"role": "system", "content": system},
                  {"role": "user", "content": user}],
        max_tokens=max_tokens,
        temperature=temperature,
    )
    return resp["choices"][0]["message"]["content"].strip()

# ── Public API ───────────────────────────────────────────────────────────────

def compile_query(query: str) -> Optional[QuantumSpec]:
    """English → QuantumSpec via the LLM, or heuristic fallback."""
    llm = _load()
    if llm is not None:
        raw = _chat(llm, _COMPILE_SYS, query, max_tokens=64, temperature=0.1)
        spec = parse_dsl(raw)
        if spec is not None:
            return spec
        log.warning("LLM produced unparseable DSL (%r); trying heuristic.", raw)
    return interpret_query(query)

def explain_results(query: str, receipt_text: str, receipt_data: dict) -> str:
    """Results → natural-language explanation via the LLM or template fallback."""
    llm = _load()
    if llm is None:
        return _template_explanation(query, receipt_data)
    user = f"Original query: {query}\n\nReceipt:\n{receipt_text}"
    return _chat(llm, _REPORT_SYS, user, max_tokens=256, temperature=0.7)

def _template_explanation(query: str, data: dict) -> str:
    """Plain-English summary derived from statistics — no LLM needed."""
    action = data.get("action", "?")
    s = data.get("statistics", {})
    if action == "RANDOM":
        return (f"Measured {data['qubits']} qubits in superposition across "
                f"{data['shots']} shots. The distribution had "
                f"{s.get('entropy_bits', 0):.2f} bits of entropy "
                f"(max {s.get('max_entropy_bits', 0):.2f}), i.e. "
                f"{s.get('uniformity', 0):.0%} uniform — these are genuinely "
                f"random bits from quantum measurement.")
    if action == "BELL":
        return (f"Created {s.get('entangled_pairs', 1)} Bell pair(s). "
                f"{s.get('correlation', 0):.2%} of shots produced perfectly "
                f"correlated outcomes between paired qubits "
                f"(ideal: 100%), confirming entanglement.")
    if action == "GROVER":
        return (f"Grover search found the marked state |{s.get('target_state')}⟩ "
                f"with {s.get('success_probability', 0):.0%} probability after "
                f"{data['extra'].get('iterations', '?')} iterations — a "
                f"quadratic speedup over the {s.get('classical_baseline', 0):.0%} "
                f"classical baseline.")
    if action == "QAOA":
        return (f"QAOA (p={s.get('depth_p', 1)}) found a MaxCut of "
                f"{s.get('best_cut', 0)}/{s.get('max_possible_cut', 0)} on the "
                f"ring graph, approximation ratio "
                f"{s.get('approximation_ratio', 0):.2f}.")
    if action == "VQE":
        return (f"VQE converged to a ground-state energy of "
                f"{s.get('vqe_converged_energy_ha', 0)} Ha for H₂ at "
                f"{data['extra'].get('bond_distance_angstrom', '?')} Å "
                f"(analytical reference {s.get('analytical_energy_ha', 0)} Ha).")
    return "No explanation available for this action."