src/llm_client.py

"""
llm_client.py
-------------
Wraps a hosted LLM (via huggingface_hub's InferenceClient, using HF's
serverless Inference Providers) to answer warehouse-operations questions,
grounded with context retrieved from the local knowledge base (simple RAG).

Design notes
------------
* Reads the HF token from the `HF_TOKEN` environment variable, which should
  be added as a Space "secret" when deployed (Settings -> Variables and
  secrets). The public demo also works without a token: it falls back to a
  deterministic, still-useful extractive answer built from the retrieved
  knowledge-base passages, so the Space never shows a broken demo.
* Serverless model availability on the free HF Inference API changes over
  time (models get gated, deprecated, or moved between providers), so
  rather than hard-depending on a single model id, we try a short list of
  candidates in order and use the first one that responds successfully.
  `LLM_MODEL_ID` (env var) is tried first if set, ahead of the built-in list.
* On failure, the *actual* exception message (not just its type) is
  surfaced back to the UI, so a broken deployment is debuggable from the
  Space itself instead of requiring log access.
"""

import os
import time
from dataclasses import dataclass, field
from typing import List, Optional

from src.retriever import KBRetriever, RetrievedDoc

# Small, widely-available instruct models known to work well on HF's free
# serverless Inference API. Tried in order; first success wins. If
# LLM_MODEL_ID is set as an env var, it is tried first, ahead of this list.
MODEL_CANDIDATES = [
    "Qwen/Qwen2.5-7B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct",
    "meta-llama/Llama-3.2-3B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "microsoft/Phi-3.5-mini-instruct",
    "HuggingFaceH4/zephyr-7b-beta",
]

_env_model = os.environ.get("LLM_MODEL_ID")
if _env_model:
    MODEL_CANDIDATES = [_env_model] + [m for m in MODEL_CANDIDATES if m != _env_model]

DEFAULT_MODEL_ID = MODEL_CANDIDATES[0]

# The token is normally expected as `HF_TOKEN`, but we also accept a few
# common alternate secret names in case the Space was set up with a
# different name. First one found wins. Add your own name here if needed.
TOKEN_ENV_VAR_CANDIDATES = [
    "HF_TOKEN",
    "Smart_Warehouse",
    "HUGGINGFACE_TOKEN",
    "HUGGINGFACEHUB_API_TOKEN",
    "HF_API_TOKEN",
    "HUGGING_FACE_HUB_TOKEN",
    "HF_ACCESS_TOKEN",
]


def _get_hf_token() -> Optional[str]:
    for var in TOKEN_ENV_VAR_CANDIDATES:
        val = os.environ.get(var)
        if val:
            return val
    return None

SYSTEM_PROMPT = (
    "You are the Smart Warehouse AI Assistant, a helpful operations copilot "
    "for a large automated distribution center (conveyors, AS/RS, AGVs/AMRs, "
    "sortation, and a WMS). Answer concisely and practically, in the tone of "
    "an experienced warehouse operations engineer. Use the provided CONTEXT "
    "when relevant, and say so plainly if the question is outside the "
    "context. Prefer short paragraphs or bullet points over long prose."
)


@dataclass
class AssistantResponse:
    answer: str
    used_llm: bool
    sources: List[RetrievedDoc]
    latency_s: float
    model_id: str
    debug_errors: List[str] = field(default_factory=list)  # non-empty only when used_llm is False due to failures


def _extractive_fallback(query: str, sources: List[RetrievedDoc]) -> str:
    """Deterministic answer used when no HF token / API call fails, so the
    Space always returns something useful instead of an error."""
    if not sources:
        return (
            "I don't have grounded context for that yet. Try asking about "
            "inventory, order status, equipment maintenance, AGV routing, "
            "picking strategy, safety incidents, or general warehouse "
            "automation concepts."
        )
    lead = sources[0]
    bullets = "\n".join(f"- **{s.title}**: {s.text}" for s in sources)
    return (
        f"(Showing retrieved knowledge instead of an LLM-generated answer -- "
        f"see the diagnostics below.)\n\n"
        f"Based on **{lead.title}**, here's the relevant information:\n\n{bullets}"
    )


def answer_query(
    query: str,
    retriever: KBRetriever,
    k: int = 2,
    max_tokens: int = 350,
) -> AssistantResponse:
    start = time.time()
    sources = retriever.retrieve(query, k=k)
    context_block = "\n\n".join(f"[{s.title}]\n{s.text}" for s in sources)

    hf_token = _get_hf_token()

    if not hf_token:
        answer = _extractive_fallback(query, sources)
        return AssistantResponse(
            answer=answer,
            used_llm=False,
            sources=sources,
            latency_s=time.time() - start,
            model_id="extractive-fallback",
            debug_errors=[
                "No HF token secret found. Checked env vars: "
                + ", ".join(TOKEN_ENV_VAR_CANDIDATES)
                + ". Set one of these as a Space secret (Settings -> Variables and secrets)."
            ],
        )

    try:
        from huggingface_hub import InferenceClient
    except ImportError as e:
        answer = _extractive_fallback(query, sources)
        return AssistantResponse(
            answer=answer, used_llm=False, sources=sources,
            latency_s=time.time() - start, model_id="extractive-fallback",
            debug_errors=[f"huggingface_hub not importable: {e}"],
        )

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"CONTEXT:\n{context_block}\n\nQUESTION: {query}"},
    ]

    errors = []
    for candidate in MODEL_CANDIDATES:
        try:
            # provider="auto" (the huggingface_hub default) lets HF's Inference
            # Providers router pick whichever backend (hf-inference, Together,
            # Fireworks, Novita, SambaNova, etc.) actually serves this specific
            # model -- hardcoding a single provider caused "model not supported
            # by provider X" errors for models hosted elsewhere.
            client = InferenceClient(model=candidate, token=hf_token, provider="auto")
            completion = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=0.3)
            text = completion.choices[0].message.content
            if text and text.strip():
                return AssistantResponse(
                    answer=text,
                    used_llm=True,
                    sources=sources,
                    latency_s=time.time() - start,
                    model_id=candidate,
                )
            errors.append(f"{candidate}: empty response")
        except Exception as e:  # noqa: BLE001 -- try the next candidate model
            errors.append(f"{candidate}: {type(e).__name__}: {e}")

    # All candidates failed -- fall back, but surface the real errors so the
    # deployment is debuggable directly from the UI.
    answer = _extractive_fallback(query, sources)
    return AssistantResponse(
        answer=answer,
        used_llm=False,
        sources=sources,
        latency_s=time.time() - start,
        model_id="extractive-fallback",
        debug_errors=errors,
    )


def test_connection(retriever: Optional[KBRetriever] = None) -> AssistantResponse:
    """Runs a single canned query through the full pipeline -- used by the
    'Test LLM connection' diagnostics button in the app."""
    retriever = retriever or KBRetriever()
    return answer_query("What is a WMS?", retriever, k=1, max_tokens=60)