from __future__ import annotations import json import math import os import re import traceback from collections import Counter from pathlib import Path from typing import Any, Optional from fastapi import FastAPI from openai import OpenAI from pydantic import BaseModel, Field from langchain_community.vectorstores import FAISS from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from qdrant_client import QdrantClient from qdrant_client.http import models as qmodels from qdrant_client.http.models import Distance, VectorParams EMBED_DIM = 384 FAISS_DIR = Path(os.getenv("FAISS_DIR", "faiss_store")) PAGE_RE = re.compile(r"\[PAGE\s+(\d+)\]", re.IGNORECASE) FIGURE_REF_RE = re.compile(r"\b(fig(?:ure)?\.?)\s*(\d+)\b", re.IGNORECASE) TABLE_REF_RE = re.compile(r"\b(tab(?:le)?\.?)\s*(\d+)\b", re.IGNORECASE) CITATION_REF_RE = re.compile(r"\[(S\d+)\]") VISUAL_TYPES = { "figure_explain", "figure_ocr", "figure_caption", "table_explain", "table_ocr", "table_caption", "page_ocr", "table_pdfplumber", "table_unstructured", "tesseract_ocr", } _VECTORSTORE: Optional[Any] = None _EMBEDDINGS: Optional[Any] = None STOPWORDS = { "the", "and", "for", "with", "that", "this", "from", "into", "your", "about", "what", "which", "where", "when", "how", "why", "does", "did", "are", "is", "was", "were", "be", "been", "being", "have", "has", "had", "can", "could", "should", "would", "may", "might", "will", "shall", "a", "an", "of", "to", "in", "on", "at", "by", "as", "it", "its", "or", "if", "but", "not", "we", "our", "you", "they", "their", "them", } DOC_GROUNDED_RE = re.compile( r"\b(document|doc|pdf|page|citation|snippet|source|context|table|figure|selected)\b", flags=re.IGNORECASE, ) DEFAULT_MODEL_ID = "llama-3.1-8b-instant" DEFAULT_MODEL_CANDIDATES = [ "llama-3.1-8b-instant", "llama-3.3-70b-versatile", "openai/gpt-oss-20b", ] class HistoryMessage(BaseModel): role: str content: str class DocumentContext(BaseModel): documentName: Optional[str] = None documentKind: Optional[str] = None documentText: Optional[str] = None class QueryRequest(BaseModel): sessionId: str message: str userId: Optional[str] = None history: list[HistoryMessage] = Field(default_factory=list) topK: int = 6 document: Optional[DocumentContext] = None class RetrievedChunk(BaseModel): id: str page: int chunkType: str text: str startOffset: Optional[int] = None endOffset: Optional[int] = None sourceDocument: Optional[str] = None score: Optional[float] = None class Citation(BaseModel): id: str page: int chunkType: str text: str startOffset: Optional[int] = None endOffset: Optional[int] = None sourceDocument: Optional[str] = None class QueryResponse(BaseModel): answer: str retrievedChunks: list[RetrievedChunk] citations: list[Citation] app = FastAPI(title="chatqna-rag-service", version="0.2.0") def get_secret(name: str, default: Optional[str] = None) -> Optional[str]: return os.getenv(name, default) def get_bool_secret(name: str, default: bool = False) -> bool: raw = (get_secret(name, str(default)) or "").strip().lower() return raw in {"1", "true", "yes", "on"} def normalize_text(text: str) -> str: return (text or "").strip() def get_embeddings() -> Any: """ Lazy-load sentence-transformers stack to keep startup memory low. """ global _EMBEDDINGS if _EMBEDDINGS is None: from langchain_huggingface import HuggingFaceEmbeddings _EMBEDDINGS = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") return _EMBEDDINGS def get_retrieval_mode() -> str: mode = (get_secret("RAG_RETRIEVAL_MODE", "semantic") or "semantic").strip().lower() return mode if mode in {"lexical", "semantic"} else "semantic" def get_rerank_enabled() -> bool: return get_bool_secret("RAG_LLM_RERANK", True) def get_rerank_candidate_k(top_k: int) -> int: cfg = int(get_secret("RAG_RERANK_CANDIDATES", "8") or "8") cfg = max(3, min(cfg, 16)) return max(int(top_k), cfg) def hf_routed_model(model_id: str) -> str: return model_id if ":" in model_id else f"{model_id}:groq" def get_model_candidates(model_id: str) -> list[str]: preferred = normalize_text(model_id) or DEFAULT_MODEL_ID raw = (get_secret("RAG_MODEL_CANDIDATES", "") or "").strip() if raw: candidates = [normalize_text(part) for part in raw.split(",")] candidates = [c for c in candidates if c] else: candidates = [preferred] + [m for m in DEFAULT_MODEL_CANDIDATES if m != preferred] deduped: list[str] = [] seen: set[str] = set() for candidate in candidates: key = candidate.lower() if key in seen: continue seen.add(key) deduped.append(candidate) return deduped or [DEFAULT_MODEL_ID] def openai_chat( client: OpenAI, model: str, messages: list[dict[str, Any]], temperature: float, max_tokens: int, ) -> tuple[str, Any]: resp = client.chat.completions.create( model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, n=1, ) content = resp.choices[0].message.content if resp and resp.choices else "" return content or "", resp def llm_chat_with_fallback( model_id: str, messages: list[dict[str, Any]], temperature: float, max_tokens: int, hf_token: Optional[str], groq_key: Optional[str], ) -> dict[str, Any]: result: dict[str, Any] = { "content": "", "primary_used": False, "raw": None, "used_model": None, "error_primary": None, "error_fallback": None, } candidate_models = get_model_candidates(model_id) if hf_token: primary_errors: list[str] = [] try: hf_client = OpenAI(base_url="https://router.huggingface.co/v1", api_key=hf_token) for candidate in candidate_models: try: routed = hf_routed_model(candidate) content, raw = openai_chat(hf_client, routed, messages, temperature, max_tokens) if normalize_text(content): result.update( { "content": content, "primary_used": True, "raw": raw, "used_model": candidate, } ) return result primary_errors.append(f"{candidate}: empty content") except Exception as exc: # pragma: no cover primary_errors.append(f"{candidate}: {type(exc).__name__}: {exc}") except Exception as exc: # pragma: no cover primary_errors.append(f"hf_client_init: {type(exc).__name__}: {exc}") result["error_primary"] = " | ".join(primary_errors[:4]) or "Primary returned empty content." else: result["error_primary"] = "Missing HUGGINGFACE_API_TOKEN" if not groq_key: result["error_fallback"] = "Missing GROQ_API_KEY (fallback not available)" return result fallback_errors: list[str] = [] try: groq_client = OpenAI(base_url="https://api.groq.com/openai/v1", api_key=groq_key) for candidate in candidate_models: try: content, raw = openai_chat(groq_client, candidate, messages, temperature, max_tokens) if normalize_text(content): result.update({"content": content, "raw": raw, "used_model": candidate}) return result fallback_errors.append(f"{candidate}: empty content") except Exception as exc: # pragma: no cover fallback_errors.append(f"{candidate}: {type(exc).__name__}: {exc}") except Exception as exc: # pragma: no cover fallback_errors.append(f"groq_client_init: {type(exc).__name__}: {exc}") result["error_fallback"] = " | ".join(fallback_errors[:4]) or "Fallback returned empty content." return result def load_faiss(embeddings: Any, path: Path = FAISS_DIR) -> Optional[FAISS]: if (path / "index.faiss").exists() and (path / "index.pkl").exists(): return FAISS.load_local(str(path), embeddings, allow_dangerous_deserialization=True) return None def init_qdrant_vectorstore() -> tuple[Optional[QdrantVectorStore], Optional[str]]: url = get_secret("QDRANT_URL") api_key = get_secret("QDRANT_API_KEY") collection_name = get_secret("QDRANT_COLLECTION", "doc_kb") if not url or not api_key: return None, "Missing QDRANT_URL or QDRANT_API_KEY" try: client = QdrantClient(url=url, api_key=api_key) try: client.get_collection(collection_name=collection_name) except Exception: client.create_collection( collection_name=collection_name, vectors_config=VectorParams(size=EMBED_DIM, distance=Distance.COSINE), ) store = QdrantVectorStore(client=client, collection_name=collection_name, embedding=get_embeddings()) return store, None except Exception as exc: return None, f"{type(exc).__name__}: {exc}" def ensure_vectorstore() -> Optional[Any]: global _VECTORSTORE if _VECTORSTORE is not None: return _VECTORSTORE backend = (get_secret("KB_BACKEND", "qdrant") or "qdrant").lower() if backend.startswith("qdrant"): store, _ = init_qdrant_vectorstore() if store is not None: _VECTORSTORE = store return _VECTORSTORE _VECTORSTORE = load_faiss(get_embeddings()) return _VECTORSTORE def qdrant_filter_for_chunk_types(types: list[str]) -> qmodels.Filter: should: list[qmodels.FieldCondition] = [] for path in ("metadata.chunk_type", "chunk_type"): should.append(qmodels.FieldCondition(key=path, match=qmodels.MatchAny(any=types))) return qmodels.Filter(should=should) def is_visual_question(question: str) -> bool: return bool(re.search(r"\b(fig|figure|diagram|chart|graph|image|table|tab)\b", question or "", flags=re.I)) def extract_ref(question: str) -> Optional[str]: m = FIGURE_REF_RE.search(question or "") if m: return f"figure {m.group(2)}" m = TABLE_REF_RE.search(question or "") if m: return f"table {m.group(2)}" return None def retrieve_docs(vectorstore: Any, query: str, k: int, prefer_visual: bool) -> list[Any]: docs_all: list[Any] = [] if prefer_visual and isinstance(vectorstore, QdrantVectorStore): try: flt = qdrant_filter_for_chunk_types(list(VISUAL_TYPES)) docs_vis = vectorstore.similarity_search(query, k=min(12, max(k, 10)), filter=flt) docs_all.extend(docs_vis) except Exception: pass try: docs_gen = vectorstore.similarity_search(query, k=max(k, 6)) docs_all.extend(docs_gen) except Exception: return [] seen: set[str] = set() uniq: list[Any] = [] for doc in docs_all: text = getattr(doc, "page_content", "") or "" key = text.strip()[:2000] if key and key not in seen: seen.add(key) uniq.append(doc) if prefer_visual and not isinstance(vectorstore, QdrantVectorStore): vis_first: list[Any] = [] rest: list[Any] = [] for doc in uniq: md = getattr(doc, "metadata", {}) or {} chunk_type = (md.get("chunk_type") or "").lower() if chunk_type in VISUAL_TYPES: vis_first.append(doc) else: rest.append(doc) uniq = vis_first + rest return uniq[: max(k, 10) if prefer_visual else k] def infer_page_from_text(text: str) -> int: match = PAGE_RE.search(text or "") if match: try: return max(1, int(match.group(1))) except ValueError: return 1 return 1 def infer_page_from_offset(text: str, offset: int) -> int: if not text: return 1 before = text[: max(0, min(len(text), offset))] matches = list(PAGE_RE.finditer(before)) if not matches: return 1 try: return max(1, int(matches[-1].group(1))) except ValueError: return 1 def dot(a: list[float], b: list[float]) -> float: return sum(x * y for x, y in zip(a, b)) def cosine_similarity(a: list[float], b: list[float]) -> float: na = math.sqrt(dot(a, a)) nb = math.sqrt(dot(b, b)) if na == 0.0 or nb == 0.0: return 0.0 return dot(a, b) / (na * nb) def tokenize_retrieval(text: str) -> list[str]: toks = re.findall(r"[A-Za-z0-9][A-Za-z0-9\-_]{1,}", text or "") return [t.lower() for t in toks if t.lower() not in STOPWORDS] def lexical_overlap_score(query: str, text: str) -> float: q_tokens = tokenize_retrieval(query) t_tokens = tokenize_retrieval(text) if not q_tokens or not t_tokens: return 0.0 q_count = Counter(q_tokens) t_count = Counter(t_tokens) overlap = sum(min(cnt, t_count.get(tok, 0)) for tok, cnt in q_count.items()) if overlap <= 0: return 0.0 # Lightweight score favoring chunks with better query coverage. return overlap / float(len(q_tokens) + 0.35 * len(t_tokens)) def split_text_with_offsets(text: str, chunk_size: int, chunk_overlap: int) -> list[dict[str, Any]]: splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) parts = splitter.split_text(text or "") chunks: list[dict[str, Any]] = [] cursor = 0 for part in parts: probe = part[:220] if not probe: continue idx = text.find(probe, max(0, cursor - chunk_overlap - 500)) if idx < 0: idx = text.find(probe) if idx < 0: continue end = min(len(text), idx + len(part)) cursor = end chunks.append( { "text": part, "startOffset": idx, "endOffset": end, "page": infer_page_from_offset(text, idx), "chunkType": "main_text", } ) return chunks def retrieve_from_uploaded_document(payload: QueryRequest) -> list[dict[str, Any]]: document = payload.document document_text = (document.documentText if document else "") or "" if not document_text.strip(): return [] chunk_size = int(get_secret("UPLOAD_CHUNK_SIZE", "1000") or "1000") chunk_overlap = int(get_secret("UPLOAD_CHUNK_OVERLAP", "200") or "200") chunks = split_text_with_offsets(document_text, chunk_size=chunk_size, chunk_overlap=chunk_overlap) if not chunks: return [] ranked: list[dict[str, Any]] = [] if get_retrieval_mode() == "semantic": embeddings = get_embeddings() query_vec = embeddings.embed_query(payload.message) doc_vecs = embeddings.embed_documents([c["text"] for c in chunks]) for chunk, vec in zip(chunks, doc_vecs): score = cosine_similarity(query_vec, vec) ranked.append({**chunk, "score": float(score)}) else: for chunk in chunks: score = lexical_overlap_score(payload.message, str(chunk.get("text") or "")) ranked.append({**chunk, "score": float(score)}) ranked.sort(key=lambda c: c["score"], reverse=True) candidate_k = get_rerank_candidate_k(payload.topK) if get_rerank_enabled() else int(payload.topK) top_k = max(1, min(candidate_k, len(ranked))) out: list[dict[str, Any]] = [] for idx, item in enumerate(ranked[:top_k], start=1): out.append( { "id": f"S{idx}", "page": int(item["page"]), "chunkType": str(item.get("chunkType") or "main_text"), "text": str(item["text"]), "startOffset": int(item["startOffset"]), "endOffset": int(item["endOffset"]), "sourceDocument": document.documentName if document else None, "score": round(float(item.get("score", 0.0)), 5), } ) return out def locate_offsets_in_document(text: str, snippet: str) -> tuple[Optional[int], Optional[int]]: source = text or "" needle = (snippet or "").strip() if not source or not needle: return None, None probe = needle[:260] idx = source.lower().find(probe.lower()) if idx < 0: return None, None return idx, min(len(source), idx + len(needle)) def retrieve_from_vectorstore(payload: QueryRequest) -> list[dict[str, Any]]: vectorstore = ensure_vectorstore() if vectorstore is None: return [] prefer_visual = is_visual_question(payload.message) requested_k = get_rerank_candidate_k(payload.topK) if get_rerank_enabled() else int(payload.topK) k = max(int(requested_k), 10) if prefer_visual else int(requested_k) docs = retrieve_docs(vectorstore, payload.message, k=k, prefer_visual=prefer_visual) if not docs: return [] document_text = ((payload.document.documentText if payload.document else "") or "").strip() source_document = payload.document.documentName if payload.document else None out: list[dict[str, Any]] = [] for idx, doc in enumerate(docs, start=1): snippet = (getattr(doc, "page_content", "") or "").strip() if not snippet: continue md = getattr(doc, "metadata", {}) or {} page = md.get("page") or md.get("page_number") or infer_page_from_text(snippet) try: page_num = max(1, int(page)) except Exception: page_num = 1 start_offset = md.get("start_offset") end_offset = md.get("end_offset") if start_offset is None or end_offset is None: start_offset, end_offset = locate_offsets_in_document(document_text, snippet) out.append( { "id": f"S{idx}", "page": page_num, "chunkType": str(md.get("chunk_type") or "main_text"), "text": snippet, "startOffset": int(start_offset) if isinstance(start_offset, int) else None, "endOffset": int(end_offset) if isinstance(end_offset, int) else None, "sourceDocument": source_document or md.get("filename"), "score": None, } ) return out def extract_json_object(text: str) -> Optional[dict[str, Any]]: raw = (text or "").strip() if not raw: return None try: data = json.loads(raw) return data if isinstance(data, dict) else None except Exception: pass fenced = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", raw, flags=re.IGNORECASE) if fenced: block = fenced.group(1).strip() try: data = json.loads(block) return data if isinstance(data, dict) else None except Exception: pass start = raw.find("{") end = raw.rfind("}") if start >= 0 and end > start: try: data = json.loads(raw[start : end + 1]) return data if isinstance(data, dict) else None except Exception: return None return None def llm_rerank_chunks( question: str, history: list[dict[str, str]], chunks: list[dict[str, Any]], model_id: str, hf_token: Optional[str], groq_key: Optional[str], ) -> tuple[list[dict[str, Any]], Optional[str]]: if not get_rerank_enabled() or len(chunks) < 2: return chunks, None candidate_limit = min(len(chunks), get_rerank_candidate_k(len(chunks))) candidates = chunks[:candidate_limit] by_id = {str(c.get("id")): c for c in candidates if c.get("id")} if len(by_id) < 2: return chunks, None history_tail = history[-4:] if history else [] history_text = "\n".join(f"{m.get('role', 'user')}: {(m.get('content') or '')[:240]}" for m in history_tail) candidate_lines: list[str] = [] for c in candidates: sid = str(c.get("id") or "") page = c.get("page") or 1 ctype = c.get("chunkType") or "main_text" txt = (c.get("text") or "").strip().replace("\n", " ") if len(txt) > 800: txt = txt[:800] + "..." candidate_lines.append(f"[{sid}] page={page} type={ctype} text={txt}") rerank_messages: list[dict[str, str]] = [ { "role": "system", "content": ( "You are a retrieval ranker. Return ONLY JSON with keys: " "ordered_ids (array of snippet IDs sorted by relevance), " "needs_clarification (boolean), clarification_question (string). " "Do not answer the user question directly. " "If the user query is ambiguous relative to snippets, set needs_clarification=true and ask one concise question." ), }, { "role": "user", "content": ( f"Conversation tail:\n{history_text or '(none)'}\n\n" f"Question:\n{question}\n\n" f"Candidate snippets:\n" + "\n".join(candidate_lines) ), }, ] rerank_max_tokens = int(get_secret("RAG_RERANK_MAX_TOKENS", "280") or "280") rerank_result = llm_chat_with_fallback( model_id=model_id, messages=rerank_messages, temperature=0.0, max_tokens=max(96, min(rerank_max_tokens, 512)), hf_token=hf_token, groq_key=groq_key, ) rerank_content = normalize_text(rerank_result.get("content", "")) parsed = extract_json_object(rerank_content) if not parsed: return chunks, None needs_clarification = bool(parsed.get("needs_clarification")) clarification_question = normalize_text(str(parsed.get("clarification_question") or "")) if needs_clarification and clarification_question: return chunks, clarification_question ordered_ids_raw = parsed.get("ordered_ids") or [] if not isinstance(ordered_ids_raw, list): return chunks, None ordered_ids: list[str] = [] for item in ordered_ids_raw: sid = str(item or "").strip() if sid and sid in by_id and sid not in ordered_ids: ordered_ids.append(sid) if not ordered_ids: return chunks, None re_ranked_candidates = [by_id[sid] for sid in ordered_ids] remaining_candidates = [c for c in candidates if str(c.get("id")) not in ordered_ids] return re_ranked_candidates + remaining_candidates + chunks[candidate_limit:], None def build_qa_prompt_with_history( history: list[dict[str, str]], context_blocks: list[str], question: str, max_history_turns: int = 8, ) -> list[dict[str, str]]: msgs = [m for m in history if m.get("role") in ("user", "assistant")] if len(msgs) > max_history_turns * 2: msgs = msgs[-max_history_turns * 2 :] messages: list[dict[str, str]] = [ { "role": "system", "content": ( "You are a precise assistant. Prefer answers grounded in the provided snippets. " "If the query is ambiguous (for example, 'explain this') and target passage is unclear, " "ask one concise clarifying question instead of guessing. " "If the answer is not present in snippets, say: 'Not found in the knowledge base.' " "When you use snippets, cite snippet IDs like [S1], [S2]. " "For TABLE questions, prioritize [TABLE_EXPLAIN], [TABLE_OCR], [TABLE_CAPTION], " "[TABLE_PDFPLUMBER], [TABLE_UNSTRUCTURED], [PAGE_OCR] snippets. " "For FIGURE questions, prioritize [FIGURE_EXPLAIN], [FIGURE_OCR], [FIGURE_CAPTION], [PAGE_OCR] snippets." ), } ] messages.extend(msgs) context_blob = "\n\n".join(context_blocks) messages.append({"role": "user", "content": f"Snippets:\n{context_blob}\n\nQuestion: {question}\nAnswer:"}) return messages def build_general_chat_prompt(history: list[dict[str, str]], question: str, max_history_turns: int = 8) -> list[dict[str, str]]: msgs = [m for m in history if m.get("role") in ("user", "assistant")] if len(msgs) > max_history_turns * 2: msgs = msgs[-max_history_turns * 2 :] messages: list[dict[str, str]] = [ { "role": "system", "content": ( "You are ChatQnA, a concise and helpful assistant. " "Answer naturally. If user asks document-specific questions without available context, " "ask them to upload/select the relevant document section." ), } ] messages.extend(msgs) messages.append({"role": "user", "content": question}) return messages def is_doc_grounded_query(question: str) -> bool: return bool(DOC_GROUNDED_RE.search(question or "")) def build_local_fallback_answer(chunks: list[dict[str, Any]]) -> str: if not chunks: return "No relevant content found in the knowledge base." lines = [] for chunk in chunks[:3]: summary = (chunk.get("text") or "").strip().replace("\n", " ") if len(summary) > 220: summary = summary[:217] + "..." lines.append(f"- {summary} [{chunk['id']}]") return "Based on retrieved context:\n" + "\n".join(lines) def build_citations(answer: str, chunks: list[dict[str, Any]]) -> list[Citation]: by_id = {str(c["id"]): c for c in chunks} ids: list[str] = [] for sid in CITATION_REF_RE.findall(answer or ""): if sid in by_id and sid not in ids: ids.append(sid) if not ids: return [] citations: list[Citation] = [] for sid in ids: chunk = by_id.get(sid) if not chunk: continue citations.append( Citation( id=sid, page=int(chunk.get("page") or 1), chunkType=str(chunk.get("chunkType") or "main_text"), text=str(chunk.get("text") or ""), startOffset=chunk.get("startOffset"), endOffset=chunk.get("endOffset"), sourceDocument=chunk.get("sourceDocument"), ) ) return citations @app.get("/health") def health() -> dict[str, str]: backend = type(_VECTORSTORE).__name__ if _VECTORSTORE is not None else "lazy_uninitialized" return {"status": "ok", "service": "rag-service", "vectorstore": backend} @app.post("/query", response_model=QueryResponse) def query(payload: QueryRequest) -> QueryResponse: try: model_id = get_secret("RAG_MODEL_ID", DEFAULT_MODEL_ID) or DEFAULT_MODEL_ID temperature = float(get_secret("RAG_TEMPERATURE", "0.2") or "0.2") max_tokens = int(get_secret("RAG_MAX_TOKENS", "512") or "512") hf_token = get_secret("HUGGINGFACE_API_TOKEN") groq_key = get_secret("GROQ_API_KEY") history = [m.model_dump() for m in payload.history] retrieved = retrieve_from_uploaded_document(payload) if not retrieved: retrieved = retrieve_from_vectorstore(payload) if not retrieved: if not is_doc_grounded_query(payload.message): general_messages = build_general_chat_prompt(history, payload.message, max_history_turns=8) general = llm_chat_with_fallback( model_id=model_id, messages=general_messages, temperature=temperature, max_tokens=max_tokens, hf_token=hf_token, groq_key=groq_key, ) general_answer = normalize_text(general.get("content", "")) if general_answer: return QueryResponse(answer=general_answer, retrievedChunks=[], citations=[]) return QueryResponse( answer=( "I can answer this, but the answer model is currently unavailable. " "Please retry shortly." ), retrievedChunks=[], citations=[], ) return QueryResponse(answer="No relevant content found in the knowledge base.", retrievedChunks=[], citations=[]) retrieved, clarifying_question = llm_rerank_chunks( question=payload.message, history=history, chunks=retrieved, model_id=model_id, hf_token=hf_token, groq_key=groq_key, ) if clarifying_question: return QueryResponse(answer=clarifying_question, retrievedChunks=[], citations=[]) retrieved = retrieved[: max(1, int(payload.topK))] context_blocks = [f"[{c['id']}] {c['text']}" for c in retrieved] messages = build_qa_prompt_with_history(history, context_blocks, payload.message, max_history_turns=8) result = llm_chat_with_fallback( model_id=model_id, messages=messages, temperature=temperature, max_tokens=max_tokens, hf_token=hf_token, groq_key=groq_key, ) answer = normalize_text(result.get("content", "")) if not answer: answer = build_local_fallback_answer(retrieved) if result.get("error_primary") or result.get("error_fallback"): print( "llm_unavailable:", { "model_candidates": get_model_candidates(model_id), "error_primary": result.get("error_primary"), "error_fallback": result.get("error_fallback"), }, ) answer += "\n\n(LLM unavailable right now; showing highest-signal retrieved context.)" if "not found in the knowledge base" in answer.lower() and context_blocks: retry_messages = [ { "role": "system", "content": "You MUST answer using the snippets below. Do NOT say 'Not found' if any relevant content exists.", }, {"role": "user", "content": f"Snippets:\n{'\n\n'.join(context_blocks)}\n\nQuestion: {payload.message}\nAnswer:"}, ] retry = llm_chat_with_fallback( model_id=model_id, messages=retry_messages, temperature=temperature, max_tokens=max_tokens, hf_token=hf_token, groq_key=groq_key, ) retry_text = normalize_text(retry.get("content", "")) if retry_text: answer = retry_text citations = build_citations(answer, retrieved) response_chunks = [RetrievedChunk(**chunk) for chunk in retrieved] return QueryResponse(answer=answer, retrievedChunks=response_chunks, citations=citations) except Exception as exc: trace = traceback.format_exc(limit=4) print("query_exception:", trace) fallback = RetrievedChunk( id="S1", page=1, chunkType="main_text", text=f"rag-service exception: {type(exc).__name__}: {exc}", ) return QueryResponse( answer=( "RAG query failed and returned a guarded fallback response. " "Check rag-service logs for details." ), retrievedChunks=[fallback], citations=[Citation(id="S1", page=1, chunkType="main_text", text=fallback.text)], )