""" 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)