infer.py

"""
Thin wrapper around a small GGUF model run via llama-cpp-python for CPU inference.

Two calls are used:
- generate_sql(question)   -> raw SQL string (unsafe until passed through guardrail.py)
- synthesize_answer(...)   -> plain-language answer from query results

Model: Qwen2.5-3B-Instruct-GGUF (Q4_K_M). Swap MODEL_REPO/MODEL_FILE for the
1.5B build if the 3B feels slow on the Spaces free CPU tier.
"""

import os
from llama_cpp import Llama

MODEL_REPO = os.environ.get("MODEL_REPO", "Qwen/Qwen2.5-3B-Instruct-GGUF")
MODEL_FILE = os.environ.get("MODEL_FILE", "qwen2.5-3b-instruct-q4_k_m.gguf")

_llm = None


def get_llm():
    global _llm
    if _llm is None:
        # from_pretrained downloads the GGUF from the Hub on first run and
        # caches it locally (works out of the box on HF Spaces).
        _llm = Llama.from_pretrained(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            n_ctx=4096,
            n_threads=os.cpu_count(),
            verbose=False,
        )
    return _llm


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


def generate_sql(prompt: str) -> str:
    """`prompt` should already be the full schema+few-shot+question text
    from llm/prompt_template.py:build_prompt(). Returns raw SQL text."""
    raw = _chat(
        system_prompt="You write SQLite SELECT queries. Output only the SQL, nothing else.",
        user_prompt=prompt,
        max_tokens=200,
        temperature=0.0,
    )
    # models sometimes wrap output in ```sql fences despite instructions — strip them
    raw = raw.replace("```sql", "").replace("```", "").strip()
    return raw


def synthesize_answer(question: str, columns: list, rows: list) -> str:
    if not rows:
        return "No matching results found."
    preview = rows[:20]
    table_text = " | ".join(columns) + "\n" + "\n".join(
        " | ".join(str(v) for v in r) for r in preview
    )
    user_prompt = (
        f"Question: {question}\n\n"
        f"Query result ({len(rows)} row(s), showing up to 20):\n{table_text}\n\n"
        "Answer the question in one or two plain sentences, referencing the "
        "actual numbers/names from the result. Do not mention SQL."
    )
    return _chat(
        system_prompt="You are a concise network operations assistant.",
        user_prompt=user_prompt,
        max_tokens=200,
        temperature=0.2,
    )