Spaces:
Sleeping
Sleeping
File size: 4,268 Bytes
f0fae3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """
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",
)
|