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.
* The model name is configurable via `LLM_MODEL_ID` (defaults to a small,
  fast, freely-hostable instruct model).
"""

import os
import time
from dataclasses import dataclass
from typing import List

from src.retriever import KBRetriever, RetrievedDoc

DEFAULT_MODEL_ID = os.environ.get("LLM_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")

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


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"(Offline / no LLM API key configured -- showing retrieved "
        f"knowledge instead of a generated answer.)\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,
    model_id: str = DEFAULT_MODEL_ID,
    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 = os.environ.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",
        )

    try:
        from huggingface_hub import InferenceClient

        client = InferenceClient(model=model_id, token=hf_token)
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": f"CONTEXT:\n{context_block}\n\nQUESTION: {query}",
            },
        ]
        completion = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=0.3)
        text = completion.choices[0].message.content
        return AssistantResponse(
            answer=text,
            used_llm=True,
            sources=sources,
            latency_s=time.time() - start,
            model_id=model_id,
        )
    except Exception as e:  # noqa: BLE001 -- deliberately broad: any API/network issue -> fallback
        answer = _extractive_fallback(query, sources)
        answer += f"\n\n_(LLM call failed: {type(e).__name__}. Showing retrieval-only answer.)_"
        return AssistantResponse(
            answer=answer,
            used_llm=False,
            sources=sources,
            latency_s=time.time() - start,
            model_id="extractive-fallback",
        )