# Dependencies (add to requirements.txt on Spaces) # gradio pypdf pypdfium2 langchain-text-splitters sentence-transformers faiss-cpu transformers requests import base64 import hmac import hashlib import io import json import os import pickle import re import textwrap from collections import Counter from datetime import datetime, timezone from pathlib import Path from threading import Lock from typing import Any, Dict, List, Optional, Tuple from urllib.parse import parse_qs, urlparse from uuid import uuid4 import faiss import gradio as gr import numpy as np import pypdfium2 as pdfium import requests try: from gradio_rangeslider import RangeSlider except ImportError: RangeSlider = None from langchain_text_splitters import RecursiveCharacterTextSplitter from pypdf import PdfReader from sentence_transformers import SentenceTransformer from starlette.middleware import Middleware # --------------------------------------------------------------------------- # Persistent storage paths # --------------------------------------------------------------------------- def _can_write_to_dir(path: Path) -> bool: try: path.mkdir(parents=True, exist_ok=True) probe = path / ".write_probe" with probe.open("w", encoding="utf-8") as fh: fh.write("") probe.unlink(missing_ok=True) return True except OSError: return False def resolve_data_dir() -> Path: configured_dir = os.getenv("PERSIST_DIR", "").strip() candidates: List[Path] = [] if configured_dir: candidates.append(Path(configured_dir)) candidates.extend( [ Path("/data"), Path.cwd() / ".data", Path("/tmp/rp-chatbot-data"), ] ) seen_paths = set() for candidate in candidates: candidate_key = str(candidate) if candidate_key in seen_paths: continue seen_paths.add(candidate_key) if _can_write_to_dir(candidate): return candidate raise RuntimeError( "No writable data directory found. Set PERSIST_DIR to a writable path." ) DATA_DIR = resolve_data_dir() UPLOAD_DIR = DATA_DIR / "uploads" INDEX_PATH_TEMPLATE = "docs_{strategy}.faiss" META_PATH_TEMPLATE = "meta_{strategy}.pkl" LEGACY_INDEX_PATH = DATA_DIR / "docs.faiss" LEGACY_META_PATH = DATA_DIR / "meta.pkl" REGISTRY_PATH = DATA_DIR / "registry.json" FEEDBACK_PATH = DATA_DIR / "feedback.json" UPLOAD_DIR.mkdir(parents=True, exist_ok=True) # Thread-safety when running on Spaces INDEX_LOCK = Lock() REGISTRY_LOCK = Lock() FEEDBACK_LOCK = Lock() EMBEDDER_LOCK = Lock() MIGRATION_LOCK = Lock() MIGRATION_DONE = False # Embedding strategies (cosine similarity via normalized embeddings) EMBEDDING_STRATEGIES: Dict[str, Dict[str, str]] = { "fast_minilm_l6": { "label": "Fast MiniLM-L6", "model": os.getenv( "EMBED_MODEL_FAST", "sentence-transformers/all-MiniLM-L6-v2" ), }, "quality_mpnet": { "label": "Quality MPNet", "model": os.getenv("EMBED_MODEL_QUALITY", "sentence-transformers/all-mpnet-base-v2"), }, "qa_multilingual": { "label": "QA Multilingual", "model": os.getenv( "EMBED_MODEL_QA", "sentence-transformers/multi-qa-mpnet-base-dot-v1" ), }, } EMBEDDING_OPTIONS = [cfg["label"] for cfg in EMBEDDING_STRATEGIES.values()] EMBEDDING_LABEL_TO_KEY = { cfg["label"]: key for key, cfg in EMBEDDING_STRATEGIES.items() } DEFAULT_EMBEDDING_KEY = next(iter(EMBEDDING_STRATEGIES.keys())) DEFAULT_EMBEDDING_LABEL = EMBEDDING_STRATEGIES[DEFAULT_EMBEDDING_KEY]["label"] EMBEDDERS: Dict[str, SentenceTransformer] = {} # Text splitter config (tweak via env if needed) CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "800")) CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "120")) MIN_CHUNK_TEXT_LENGTH = int(os.getenv("MIN_CHUNK_TEXT_LENGTH", "80")) ACCESS_QUERY_PARAM = os.getenv("APP_ACCESS_QUERY_PARAM", "token") URL_ACCESS_TOKEN = os.getenv("APP_URL_TOKEN", "").strip() ACCESS_COOKIE_NAME = os.getenv("APP_ACCESS_COOKIE_NAME", "rag_access") ALLOW_LOOPBACK_BYPASS = os.getenv("APP_ALLOW_LOOPBACK_BYPASS", "1").strip().lower() in { "1", "true", "yes", "on", } IS_HF_SPACE = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID")) ENABLE_TOKEN_GATE_ON_SPACES = os.getenv( "APP_ENABLE_TOKEN_GATE_ON_SPACES", "0" ).strip().lower() in {"1", "true", "yes", "on"} TOKEN_GATE_ENABLED = bool(URL_ACCESS_TOKEN) and ( not IS_HF_SPACE or ENABLE_TOKEN_GATE_ON_SPACES ) def _get_env_float(name: str, default: float) -> float: try: return float(os.getenv(name, default)) except (TypeError, ValueError): return default RAG_MIN_SCORE = _get_env_float("RAG_MIN_SCORE", 0.15) def utc_now() -> datetime: return datetime.now(timezone.utc) def utc_year() -> int: return utc_now().year def utc_now_iso_z() -> str: return utc_now().isoformat(timespec="seconds").replace("+00:00", "Z") DEFAULT_CHAT_MODEL = os.getenv( "DEEPINFRA_RAG_MODEL", "meta-llama/Meta-Llama-3.1-8B-Instruct" ) REASONING_PRESETS = { "Fast (8B)": { "model": os.getenv( "DEEPINFRA_FAST_MODEL", "meta-llama/Meta-Llama-3.1-8B-Instruct" ), "temperature": _get_env_float("DEEPINFRA_FAST_TEMP", 0.1), "extra": {}, }, "Balanced (70B)": { "model": os.getenv( "DEEPINFRA_BALANCED_MODEL", "meta-llama/Meta-Llama-3.1-70B-Instruct" ), "temperature": _get_env_float("DEEPINFRA_BALANCED_TEMP", 0.15), "extra": {}, }, "Deep Reasoning (R1 Distill)": { "model": os.getenv( "DEEPINFRA_REASONING_MODEL", "deepseek-ai/DeepSeek-R1-Distill-Llama-70B" ), "temperature": _get_env_float("DEEPINFRA_REASONING_TEMP", 0.2), "extra": {"reasoning": {"effort": "medium"}}, }, } # --------------------------------------------------------------------------- # Registry helpers # --------------------------------------------------------------------------- def _load_registry() -> Dict: if not REGISTRY_PATH.exists(): return {"documents": {}} with REGISTRY_PATH.open("r", encoding="utf-8") as fh: return json.load(fh) def _save_registry(data: Dict) -> None: with REGISTRY_PATH.open("w", encoding="utf-8") as fh: json.dump(data, fh, indent=2, ensure_ascii=True) def load_registry_threadsafe() -> Dict: with REGISTRY_LOCK: return _load_registry() def update_registry(doc_id: str, payload: Dict) -> None: with REGISTRY_LOCK: registry = _load_registry() registry.setdefault("documents", {}) registry["documents"][doc_id] = payload _save_registry(registry) def remove_from_registry(doc_id: str) -> Optional[Dict]: with REGISTRY_LOCK: registry = _load_registry() documents = registry.get("documents", {}) document = documents.pop(doc_id, None) registry["documents"] = documents _save_registry(registry) return document def _load_feedback_registry() -> Dict: if not FEEDBACK_PATH.exists(): return {"feedback": {}} with FEEDBACK_PATH.open("r", encoding="utf-8") as fh: data = json.load(fh) if not isinstance(data, dict): return {"feedback": {}} feedback = data.get("feedback") if not isinstance(feedback, dict): data["feedback"] = {} return data def _save_feedback_registry(data: Dict) -> None: with FEEDBACK_PATH.open("w", encoding="utf-8") as fh: json.dump(data, fh, indent=2, ensure_ascii=True) def load_feedback_registry_threadsafe() -> Dict: with FEEDBACK_LOCK: return _load_feedback_registry() def create_feedback_entry(payload: Dict) -> Dict: with FEEDBACK_LOCK: store = _load_feedback_registry() feedback_map = store.setdefault("feedback", {}) feedback_id = payload.get("feedback_id") if not feedback_id: feedback_id = f"fb_{utc_now().strftime('%Y%m%d%H%M%S')}_{uuid4().hex[:8]}" payload["feedback_id"] = feedback_id feedback_map[feedback_id] = payload _save_feedback_registry(store) return payload def get_feedback_entry(feedback_id: str) -> Optional[Dict]: feedback_id = (feedback_id or "").strip() if not feedback_id: return None store = load_feedback_registry_threadsafe() feedback_map = store.get("feedback", {}) entry = feedback_map.get(feedback_id) return dict(entry) if isinstance(entry, dict) else None def delete_feedback_entry(feedback_id: str) -> Dict: feedback_id = (feedback_id or "").strip() if not feedback_id: return {"status": "error", "message": "Feedback ID is missing."} with FEEDBACK_LOCK: store = _load_feedback_registry() feedback_map = store.get("feedback", {}) removed = feedback_map.pop(feedback_id, None) store["feedback"] = feedback_map _save_feedback_registry(store) if not removed: return {"status": "error", "message": f"Feedback {feedback_id} not found."} return {"status": "deleted", "message": f"Deleted feedback {feedback_id}."} def mark_feedback_resolved(feedback_id: str) -> Dict: feedback_id = (feedback_id or "").strip() if not feedback_id: return {"status": "error", "message": "Feedback ID is missing."} with FEEDBACK_LOCK: store = _load_feedback_registry() feedback_map = store.get("feedback", {}) entry = feedback_map.get(feedback_id) if not isinstance(entry, dict): return {"status": "error", "message": f"Feedback {feedback_id} not found."} if entry.get("status") == "resolved": return {"status": "noop", "message": f"Feedback {feedback_id} is already resolved."} entry["status"] = "resolved" entry["resolved_at"] = utc_now_iso_z() feedback_map[feedback_id] = entry store["feedback"] = feedback_map _save_feedback_registry(store) return {"status": "resolved", "message": f"Marked feedback {feedback_id} as resolved."} def compute_sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as fh: for block in iter(lambda: fh.read(8192), b""): digest.update(block) return digest.hexdigest() # --------------------------------------------------------------------------- # Vector index helpers # --------------------------------------------------------------------------- def strategy_paths(strategy_key: str) -> Tuple[Path, Path]: index_path = DATA_DIR / INDEX_PATH_TEMPLATE.format(strategy=strategy_key) meta_path = DATA_DIR / META_PATH_TEMPLATE.format(strategy=strategy_key) return index_path, meta_path def load_embedder(strategy_key: str) -> SentenceTransformer: with EMBEDDER_LOCK: existing = EMBEDDERS.get(strategy_key) if existing is not None: return existing strategy = EMBEDDING_STRATEGIES.get(strategy_key) if not strategy: raise ValueError(f"Unknown embedding strategy: {strategy_key}") model_id = strategy["model"] embedder = SentenceTransformer(model_id) EMBEDDERS[strategy_key] = embedder return embedder def encode_texts(texts: List[str], strategy_key: str) -> np.ndarray: embedder = load_embedder(strategy_key) vectors = embedder.encode(texts, normalize_embeddings=True) return np.asarray(vectors, dtype="float32") def load_index(strategy_key: str) -> Tuple[Optional[faiss.Index], List[Dict]]: index_path, meta_path = strategy_paths(strategy_key) if not index_path.exists() or not meta_path.exists(): return None, [] index = faiss.read_index(str(index_path)) with meta_path.open("rb") as fh: metadata = pickle.load(fh) return index, metadata def save_index(strategy_key: str, index: faiss.Index, metadata: List[Dict]) -> None: index_path, meta_path = strategy_paths(strategy_key) faiss.write_index(index, str(index_path)) with meta_path.open("wb") as fh: pickle.dump(metadata, fh) def clear_legacy_index_files() -> None: for path in (LEGACY_INDEX_PATH, LEGACY_META_PATH): try: path.unlink() except FileNotFoundError: continue def get_registry_doc_ids() -> set[str]: registry = load_registry_threadsafe() documents = registry.get("documents", {}) return set(documents.keys()) def ensure_strategy_indexes_initialized() -> None: """ Migrates legacy single-index storage into all configured embedding strategies. """ global MIGRATION_DONE with MIGRATION_LOCK: if MIGRATION_DONE: return MIGRATION_DONE = True if not (LEGACY_INDEX_PATH.exists() and LEGACY_META_PATH.exists()): return has_new_layout = False for strategy_key in EMBEDDING_STRATEGIES: index_path, meta_path = strategy_paths(strategy_key) if index_path.exists() and meta_path.exists(): has_new_layout = True break if has_new_layout: return try: with LEGACY_META_PATH.open("rb") as fh: legacy_metadata = pickle.load(fh) except Exception: return if not isinstance(legacy_metadata, list) or not legacy_metadata: clear_legacy_index_files() return active_doc_ids = get_registry_doc_ids() if active_doc_ids: legacy_metadata = [ item for item in legacy_metadata if item.get("doc_id") in active_doc_ids ] else: legacy_metadata = [] if not legacy_metadata: clear_index_files() clear_legacy_index_files() return for item in legacy_metadata: try: int(item.get("year")) except (TypeError, ValueError): item["year"] = get_registry_doc_year(item.get("doc_id")) or utc_year() for strategy_key in EMBEDDING_STRATEGIES: strategy_metadata = [dict(item) for item in legacy_metadata] rebuild_index_from_metadata(strategy_key, strategy_metadata) clear_legacy_index_files() # --------------------------------------------------------------------------- # OCR client for DeepInfra DeepSeek OCR # --------------------------------------------------------------------------- class DeepInfraOCR: def __init__(self) -> None: self.api_token = os.getenv("DEEPINFRA") self.model_id = os.getenv("DEEPINFRA_OCR_MODEL", "deepseek-ai/deepseek-ocr") base_url = os.getenv("DEEPINFRA_API_URL", "https://api.deepinfra.com/v1/inference") self.endpoint = base_url.rstrip("/") + f"/{self.model_id}" self.timeout = int(os.getenv("DEEPINFRA_TIMEOUT", "120")) def available(self) -> bool: return bool(self.api_token) def infer(self, image) -> str: """ Runs OCR on a single PIL image. Expects DeepInfra API to respond with JSON containing the recognized text. """ if not self.available(): raise RuntimeError("DeepInfra token not configured (DEEPINFRA).") buffer = io.BytesIO() image.save(buffer, format="PNG") image_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8") payload = {"input": {"image": image_base64}} headers = { "Authorization": f"Bearer {self.api_token}", "Content-Type": "application/json", } response = requests.post( self.endpoint, headers=headers, json=payload, timeout=self.timeout, ) response.raise_for_status() data = response.json() text = self._extract_text(data) if not text: raise ValueError(f"DeepInfra OCR returned an empty result: {data}") return text.strip() @staticmethod def _extract_text(data: Dict) -> Optional[str]: """ DeepInfra responses may vary; try a few common response shapes. """ if data is None: return None if isinstance(data, str): return data for key in ("text", "output", "result"): value = data.get(key) if isinstance(value, str): return value if isinstance(value, dict): nested = ( value.get("text") or value.get("output") or value.get("result") or value.get("content") ) if isinstance(nested, str): return nested results = data.get("results") or data.get("outputs") or data.get("choices") if isinstance(results, list) and results: first = results[0] if isinstance(first, str): return first if isinstance(first, dict): return ( first.get("text") or first.get("output") or first.get("result") or first.get("content") or first.get("message", {}).get("content") ) return None OCR_CLIENT = DeepInfraOCR() # --------------------------------------------------------------------------- # PDF text extraction & chunking # --------------------------------------------------------------------------- def extract_pdf_text(pdf_path: Path) -> List[Dict]: """ Returns a list of {"page": int, "text": str} dictionaries. Falls back to OCR for pages where text extraction fails. """ reader = PdfReader(str(pdf_path)) pages: List[Dict] = [] ocr_candidates: List[int] = [] for idx, page in enumerate(reader.pages): raw_text = (page.extract_text() or "").strip() if len(raw_text) > 30: pages.append({"page": idx + 1, "text": raw_text}) else: pages.append({"page": idx + 1, "text": ""}) ocr_candidates.append(idx) if ocr_candidates and not OCR_CLIENT.available(): # Leave blank text; user will be notified about missing OCR capability. return pages if ocr_candidates: pdf = pdfium.PdfDocument(str(pdf_path)) for page_index in ocr_candidates: page = pdf.get_page(page_index) bitmap = page.render(scale=2.0, rotation=0) image = bitmap.to_pil() try: text = OCR_CLIENT.infer(image) except Exception as exc: pages[page_index]["text"] = "" pages[page_index]["error"] = str(exc) else: pages[page_index]["text"] = text return pages def extract_candidate_years(text: str) -> List[int]: current_year = utc_year() max_year = current_year + 1 years: List[int] = [] for match in re.findall(r"\b(19\d{2}|20\d{2}|21\d{2})\b", text or ""): year = int(match) if 1900 <= year <= max_year: years.append(year) return years def infer_publication_year(pdf_path: Path, doc_name: str, pages: List[Dict]) -> Tuple[int, str]: """ Extract publication year from metadata and front matter first. If unclear, use the most recent year mentioned in the document as a proxy. """ current_year = utc_year() metadata_years: List[int] = [] try: reader = PdfReader(str(pdf_path)) metadata = reader.metadata or {} for key, value in metadata.items(): if value is None: continue combined = f"{key} {value}" metadata_years.extend(extract_candidate_years(str(combined))) except Exception: pass filename_years = extract_candidate_years(doc_name) cover_text = "\n".join((p.get("text") or "") for p in pages[:3]) front_matter_years = extract_candidate_years(cover_text) explicit_years = metadata_years + filename_years + front_matter_years if explicit_years: year_counts = Counter(explicit_years) highest_count = max(year_counts.values()) tied_years = [year for year, count in year_counts.items() if count == highest_count] return max(tied_years), "explicit" full_text = "\n".join((p.get("text") or "") for p in pages) proxy_years = extract_candidate_years(full_text) if proxy_years: return max(proxy_years), "proxy_latest_mentioned" return current_year, "proxy_current_year" def chunk_pages( pages: List[Dict], doc_id: str, doc_name: str, doc_year: int ) -> List[Dict]: splitter = RecursiveCharacterTextSplitter( chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP ) chunks: List[Dict] = [] for page in pages: text = page.get("text", "").strip() if not text: continue page_chunks = splitter.split_text(text) for idx, chunk_text in enumerate(page_chunks, start=1): normalized = " ".join(chunk_text.split()) if len(normalized) < MIN_CHUNK_TEXT_LENGTH: continue chunks.append( { "text": normalized, "metadata": { "doc_id": doc_id, "doc_name": doc_name, "year": doc_year, "page": page["page"], "chunk_number": idx, }, } ) return chunks def add_chunks_to_indices(chunks: List[Dict]) -> Tuple[int, Dict[str, int]]: if not chunks: return 0, {} texts = [chunk["text"] for chunk in chunks] vectors_by_strategy: Dict[str, np.ndarray] = {} for strategy_key in EMBEDDING_STRATEGIES: vectors_by_strategy[strategy_key] = encode_texts(texts, strategy_key) with INDEX_LOCK: end_ids: Dict[str, int] = {} for strategy_key, vectors in vectors_by_strategy.items(): index, metadata = load_index(strategy_key) dim = vectors.shape[1] if index is None: index = faiss.IndexFlatIP(dim) metadata = [] elif index.d != dim: existing_texts = [item["text"] for item in metadata] index = faiss.IndexFlatIP(dim) if existing_texts: reencoded = encode_texts(existing_texts, strategy_key) index.add(reencoded) for row_idx, item in enumerate(metadata): item["vector_id"] = row_idx start_idx = len(metadata) index.add(vectors) for offset, chunk in enumerate(chunks): chunk_meta = { "vector_id": start_idx + offset, "text": chunk["text"], } chunk_meta.update(chunk["metadata"]) metadata.append(chunk_meta) save_index(strategy_key, index, metadata) end_ids[strategy_key] = start_idx + len(chunks) added = len(chunks) return added, end_ids # --------------------------------------------------------------------------- # Ingestion workflow # --------------------------------------------------------------------------- def ingest_pdf(path_str: str) -> Dict: ensure_strategy_indexes_initialized() pdf_path = Path(path_str) if not pdf_path.exists(): return {"status": "error", "message": f"File not found: {pdf_path}"} sha256 = compute_sha256(pdf_path) doc_id = sha256[:12] original_name = pdf_path.name registry = load_registry_threadsafe() existing = registry.get("documents", {}).get(doc_id) if existing: return { "status": "skipped", "doc_id": doc_id, "message": f"Already indexed ({existing['title']})", } persisted_path = UPLOAD_DIR / f"{doc_id}.pdf" if not persisted_path.exists(): persisted_path.write_bytes(pdf_path.read_bytes()) pages = extract_pdf_text(persisted_path) blank_pages = [p["page"] for p in pages if not p.get("text")] publication_year, year_source = infer_publication_year( persisted_path, original_name, pages ) chunks = chunk_pages( pages, doc_id=doc_id, doc_name=original_name, doc_year=publication_year ) added_chunks, end_ids = add_chunks_to_indices(chunks) document_record = { "doc_id": doc_id, "sha256": sha256, "title": original_name, "stored_path": str(persisted_path.relative_to(DATA_DIR)), "year": publication_year, "year_source": year_source, "pages": len(pages), "blank_pages": blank_pages, "chunks": added_chunks, "first_vector_ids": { key: (end_ids[key] - added_chunks) for key in end_ids }, "ingested_at": utc_now_iso_z(), } update_registry(doc_id, document_record) if blank_pages and OCR_CLIENT.available(): message = ( f"Indexed {original_name} ({publication_year}, {added_chunks} chunks). " f"OCR ran on pages: {blank_pages}" ) elif blank_pages: message = ( f"Indexed {original_name} ({publication_year}, {added_chunks} chunks) but " f"{len(blank_pages)} pages lacked text (set the DEEPINFRA token to enable OCR)." ) else: message = f"Indexed {original_name} ({publication_year}, {added_chunks} chunks)." return {"status": "indexed", "doc_id": doc_id, "message": message} def clear_index_files(strategy_key: Optional[str] = None) -> None: strategy_keys = [strategy_key] if strategy_key else list(EMBEDDING_STRATEGIES.keys()) for key in strategy_keys: index_path, meta_path = strategy_paths(key) for path in (index_path, meta_path): try: path.unlink() except FileNotFoundError: continue def rebuild_index_from_metadata(strategy_key: str, metadata: List[Dict]) -> None: if not metadata: clear_index_files(strategy_key=strategy_key) return texts = [m["text"] for m in metadata] vectors = encode_texts(texts, strategy_key) dim = vectors.shape[1] index = faiss.IndexFlatIP(dim) index.add(vectors) for idx, meta in enumerate(metadata): meta["vector_id"] = idx save_index(strategy_key, index, metadata) def rebuild_index_after_delete( strategy_key: str, index: faiss.Index, metadata: List[Dict], removed_doc_id: str ) -> None: if index is None or not metadata: clear_index_files(strategy_key=strategy_key) return total_vectors = int(getattr(index, "ntotal", 0)) usable_count = min(len(metadata), total_vectors) if usable_count <= 0: clear_index_files(strategy_key=strategy_key) return remaining_metadata: List[Dict] = [] remaining_vectors: List[np.ndarray] = [] for vector_idx in range(usable_count): item = metadata[vector_idx] if item.get("doc_id") == removed_doc_id: continue remaining_metadata.append(dict(item)) vector = index.reconstruct(vector_idx) remaining_vectors.append(np.asarray(vector, dtype="float32")) if not remaining_metadata: clear_index_files(strategy_key=strategy_key) return matrix = np.vstack(remaining_vectors).astype("float32") dim = matrix.shape[1] rebuilt_index = faiss.IndexFlatIP(dim) rebuilt_index.add(matrix) for new_idx, item in enumerate(remaining_metadata): item["vector_id"] = new_idx save_index(strategy_key, rebuilt_index, remaining_metadata) def delete_document(doc_id: str) -> Dict: ensure_strategy_indexes_initialized() document = remove_from_registry(doc_id) if not document: return {"status": "error", "doc_id": doc_id, "message": "Document not found."} title = document.get("title", doc_id) try: with INDEX_LOCK: for strategy_key in EMBEDDING_STRATEGIES: index, metadata = load_index(strategy_key) if not metadata or index is None: clear_index_files(strategy_key=strategy_key) continue rebuild_index_after_delete( strategy_key=strategy_key, index=index, metadata=metadata, removed_doc_id=doc_id, ) if not get_registry_doc_ids(): clear_index_files() clear_legacy_index_files() except Exception as exc: # Keep registry and index state consistent when deletion fails. update_registry(doc_id, document) return { "status": "error", "doc_id": doc_id, "message": f"Failed to remove {title}: {exc}", } stored_path = document.get("stored_path") if stored_path: stored_file = DATA_DIR / stored_path if stored_file.exists(): try: stored_file.unlink() except Exception: pass message = f"Removed {title}." return {"status": "deleted", "doc_id": doc_id, "message": message} def get_registry_doc_year(doc_id: Optional[str]) -> Optional[int]: if not doc_id: return None registry = load_registry_threadsafe() doc = registry.get("documents", {}).get(doc_id, {}) year_value = doc.get("year") try: return int(year_value) except (TypeError, ValueError): return None def resolve_hit_year(hit: Dict) -> int: try: return int(hit.get("year")) except (TypeError, ValueError): pass registry_year = get_registry_doc_year(hit.get("doc_id")) if registry_year is not None: hit["year"] = registry_year return registry_year fallback = utc_year() hit["year"] = fallback return fallback def select_diverse_hits(hits: List[Dict], top_k: int) -> List[Dict]: if len(hits) <= top_k: return hits per_doc_cap = max(2, top_k // 2) selected: List[Dict] = [] selected_ids = set() per_doc_counts: Dict[str, int] = {} for hit in hits: hit_id = ( hit.get("doc_id"), hit.get("page"), hit.get("chunk_number"), hit.get("vector_id"), ) if hit_id in selected_ids: continue doc_id = hit.get("doc_id", "") if per_doc_counts.get(doc_id, 0) >= per_doc_cap: continue selected.append(hit) selected_ids.add(hit_id) per_doc_counts[doc_id] = per_doc_counts.get(doc_id, 0) + 1 if len(selected) >= top_k: return selected for hit in hits: if len(selected) >= top_k: break hit_id = ( hit.get("doc_id"), hit.get("page"), hit.get("chunk_number"), hit.get("vector_id"), ) if hit_id in selected_ids: continue selected.append(hit) selected_ids.add(hit_id) return selected def _parse_year_range(year_from: Optional[int], year_to: Optional[int]) -> Tuple[int, int]: current_year = utc_year() from_year = int(year_from) if year_from is not None else 1900 to_year = int(year_to) if year_to is not None else current_year + 1 if from_year > to_year: from_year, to_year = to_year, from_year return from_year, to_year def parse_year_range_input(year_range: Any) -> Tuple[int, int]: year_from: Optional[int] = None year_to: Optional[int] = None if isinstance(year_range, (list, tuple)): if len(year_range) >= 1: try: year_from = int(year_range[0]) except (TypeError, ValueError): year_from = None if len(year_range) >= 2: try: year_to = int(year_range[1]) except (TypeError, ValueError): year_to = None elif year_from is not None: year_to = year_from elif year_range is not None: try: year_from = int(year_range) year_to = year_from except (TypeError, ValueError): year_from, year_to = None, None return _parse_year_range(year_from, year_to) # --------------------------------------------------------------------------- # Retrieval & LLM answer generation # --------------------------------------------------------------------------- def resolve_embedding_selection(selected_label: Optional[str]) -> Tuple[str, str]: strategy_key = EMBEDDING_LABEL_TO_KEY.get( (selected_label or "").strip(), DEFAULT_EMBEDDING_KEY ) strategy_label = EMBEDDING_STRATEGIES[strategy_key]["label"] return strategy_key, strategy_label def search_index( query: str, top_k: int = 5, strategy_key: str = DEFAULT_EMBEDDING_KEY, year_from: Optional[int] = None, year_to: Optional[int] = None, ) -> List[Dict]: ensure_strategy_indexes_initialized() active_doc_ids = get_registry_doc_ids() if not active_doc_ids: return [] from_year, to_year = _parse_year_range(year_from, year_to) with INDEX_LOCK: index, metadata = load_index(strategy_key) if index is None or not metadata: return [] query_vector = encode_texts([query], strategy_key) fetch_k = min(len(metadata), max(top_k * 8, 100)) if fetch_k < top_k: fetch_k = min(len(metadata), top_k) scores, ids = index.search(query_vector, fetch_k) candidates: List[Dict] = [] for score, idx in zip(scores[0], ids[0]): if idx == -1: continue meta = metadata[idx] doc_id = meta.get("doc_id") if doc_id not in active_doc_ids: continue hit = {**meta, "score": float(score)} year_value = resolve_hit_year(hit) if from_year <= year_value <= to_year: candidates.append(hit) if not candidates: return [] ranked = sorted(candidates, key=lambda item: item.get("score", 0.0), reverse=True) high_confidence = [item for item in ranked if item.get("score", 0.0) >= RAG_MIN_SCORE] retrieval_pool = high_confidence if high_confidence else ranked return select_diverse_hits(retrieval_pool, top_k)[:top_k] class DeepInfraChat: def __init__(self) -> None: self.api_token = os.getenv("DEEPINFRA") self.model_id = DEFAULT_CHAT_MODEL self.endpoint = os.getenv( "DEEPINFRA_CHAT_URL", "https://api.deepinfra.com/v1/openai/chat/completions" ) self.timeout = int(os.getenv("DEEPINFRA_TIMEOUT", "120")) def available(self) -> bool: return bool(self.api_token) def generate( self, messages: List[Dict], temperature: float = 0.1, model_override: Optional[str] = None, extra_payload: Optional[Dict] = None, ) -> str: if not self.available(): raise RuntimeError("Configure DEEPINFRA to enable answering.") payload = { "model": model_override or self.model_id, "messages": messages, "temperature": temperature, } if extra_payload: payload.update(extra_payload) headers = { "Authorization": f"Bearer {self.api_token}", "Content-Type": "application/json", } response = requests.post( self.endpoint, headers=headers, json=payload, timeout=self.timeout ) response.raise_for_status() data = response.json() choices = data.get("choices") if not choices: raise ValueError(f"DeepInfra chat returned no choices: {data}") content = choices[0].get("message", {}).get("content") if not content: raise ValueError(f"No content in response: {data}") return content.strip() CHAT_CLIENT = DeepInfraChat() def _token_matches(candidate: str) -> bool: if not URL_ACCESS_TOKEN: return False if not candidate: return False return hmac.compare_digest(candidate.strip(), URL_ACCESS_TOKEN) def _query_tokens_from_raw_query(raw_query: str) -> List[str]: tokens: List[str] = [] parsed = parse_qs(raw_query, keep_blank_values=True) tokens.extend(token for token in parsed.get(ACCESS_QUERY_PARAM, []) if token) # Keep raw values too so '+' and similar characters do not break matching. prefix = f"{ACCESS_QUERY_PARAM}=" for pair in raw_query.split("&"): if pair.startswith(prefix): raw_value = pair[len(prefix) :].strip() if raw_value: tokens.append(raw_value) return tokens def _query_tokens_from_scope(scope: Dict[str, Any]) -> List[str]: raw_query = (scope.get("query_string") or b"").decode("utf-8", errors="ignore") if not raw_query: return [] return _query_tokens_from_raw_query(raw_query) def _referer_query_tokens(scope: Dict[str, Any]) -> List[str]: for key, value in scope.get("headers", []): if key == b"referer": referer = value.decode("latin-1") parsed = urlparse(referer) return _query_tokens_from_raw_query(parsed.query) return [] def _cookies_from_scope(scope: Dict[str, Any]) -> Dict[str, str]: cookies: Dict[str, str] = {} for key, value in scope.get("headers", []): if key == b"cookie": header = value.decode("latin-1") for chunk in header.split(";"): if "=" not in chunk: continue cookie_key, cookie_value = chunk.strip().split("=", 1) cookies[cookie_key] = cookie_value return cookies def _decode_cookie_token(cookie_value: str) -> List[str]: if not cookie_value: return [] tokens = [cookie_value] if not cookie_value.startswith("b64."): return tokens encoded = cookie_value[4:] padded = encoded + ("=" * (-len(encoded) % 4)) try: decoded = base64.urlsafe_b64decode(padded.encode("ascii")).decode("utf-8") except Exception: return tokens if decoded: tokens.append(decoded) return tokens def _is_loopback_client(scope: Dict[str, Any]) -> bool: client = scope.get("client") if isinstance(client, (tuple, list)) and client: host = str(client[0]).lower() if host in {"127.0.0.1", "::1", "localhost"}: return True for key, value in scope.get("headers", []): if key == b"host": host_header = value.decode("latin-1").lower() if host_header.startswith("localhost:") or host_header.startswith("127.0.0.1:"): return True return False class QueryTokenGateMiddleware: def __init__(self, app) -> None: self.app = app async def __call__(self, scope, receive, send): scope_type = scope.get("type") if scope_type not in {"http", "websocket"}: await self.app(scope, receive, send) return if ALLOW_LOOPBACK_BYPASS and _is_loopback_client(scope): await self.app(scope, receive, send) return if not URL_ACCESS_TOKEN: await self.app(scope, receive, send) return # Gate only the landing page to avoid blocking internal Gradio assets/API # requests on Spaces, which can lead to a blank screen. path = str(scope.get("path") or "/") if path not in {"", "/"}: await self.app(scope, receive, send) return query_tokens = _query_tokens_from_scope(scope) referer_tokens = _referer_query_tokens(scope) cookies = _cookies_from_scope(scope) cookie_tokens = _decode_cookie_token(cookies.get(ACCESS_COOKIE_NAME, "")) token_candidates = query_tokens + referer_tokens + cookie_tokens token_valid = any(_token_matches(candidate) for candidate in token_candidates) if not token_valid: if scope_type == "websocket": await send({"type": "websocket.close", "code": 1008}) return body = f""" Token Required

Token is required

Open this Space with a valid token in the URL query:

?{ACCESS_QUERY_PARAM}=YOUR_TOKEN

""".encode("utf-8") await send( { "type": "http.response.start", "status": 200, "headers": [ (b"content-type", b"text/html; charset=utf-8"), (b"cache-control", b"no-store"), (b"content-length", str(len(body)).encode("utf-8")), ], } ) await send({"type": "http.response.body", "body": body}) return should_set_cookie = scope_type == "http" and any( _token_matches(token) for token in query_tokens ) if not should_set_cookie: await self.app(scope, receive, send) return async def send_wrapper(message): if message.get("type") == "http.response.start": headers = list(message.get("headers") or []) secure_part = "; Secure" if IS_HF_SPACE else "" cookie_token = base64.urlsafe_b64encode( URL_ACCESS_TOKEN.encode("utf-8") ).decode("ascii") cookie_value = ( f"{ACCESS_COOKIE_NAME}=b64.{cookie_token}; Path=/; HttpOnly; SameSite=Lax" f"{secure_part}" ) headers.append((b"set-cookie", cookie_value.encode("utf-8"))) message["headers"] = headers await send(message) await self.app(scope, receive, send_wrapper) def build_context_prompt(hits: List[Dict]) -> str: blocks = [] for idx, hit in enumerate(hits, start=1): snippet = hit["text"].strip().replace("\n", " ") snippet = " ".join(snippet.split()) snippet = textwrap.shorten(snippet, width=700, placeholder="…") doc_name = hit.get("doc_name", "Document") page = hit.get("page", "?") year = hit.get("year", "?") blocks.append(f"[{idx}] {doc_name} ({year}, page {page}): {snippet}") return "\n".join(blocks) def format_references(hits: List[Dict]) -> str: references_lines = [] for idx, hit in enumerate(hits, start=1): doc_name = hit.get("doc_name", "Document") page = hit.get("page", "?") year = hit.get("year", "?") score = hit.get("score", 0.0) snippet = textwrap.shorten(hit["text"].strip(), width=220, placeholder="…") references_lines.append( f"**[{idx}] {doc_name} ({year}) — page {page} (score {score:.3f})**\n{snippet}" ) return "\n\n".join(references_lines) def resolve_reasoning_selection( preset_label: Optional[str], custom_model: Optional[str] ) -> Tuple[str, float, Dict]: preset = REASONING_PRESETS.get(preset_label or "") model_id = (custom_model or "").strip() if not model_id and preset: model_id = preset.get("model", DEFAULT_CHAT_MODEL) if not model_id: model_id = DEFAULT_CHAT_MODEL temperature = preset.get("temperature", 0.1) if preset else 0.1 extra_payload = preset.get("extra", {}).copy() if preset else {} return model_id, temperature, extra_payload def tuples_to_messages(history: List[Tuple[str, str]]) -> List[Dict[str, str]]: messages: List[Dict[str, str]] = [] for user_text, bot_text in history: messages.append({"role": "user", "content": user_text}) messages.append({"role": "assistant", "content": bot_text}) return messages def condense_hits_for_trace(hits: List[Dict]) -> List[Dict]: condensed: List[Dict] = [] for hit in hits: condensed.append( { "doc_id": hit.get("doc_id"), "doc_name": hit.get("doc_name", "Document"), "year": hit.get("year"), "page": hit.get("page"), "chunk_number": hit.get("chunk_number"), "score": round(float(hit.get("score", 0.0)), 6), "snippet": textwrap.shorten( " ".join(str(hit.get("text", "")).split()), width=240, placeholder="…", ), } ) return condensed def chat_with_rag( user_message: str, history: Optional[List[Tuple[str, str]]], top_k: int, year_range: Any, embedding_label: Optional[str], preset_label: Optional[str], custom_model: Optional[str], chat_trace: Optional[List[Dict]], ) -> Tuple[List[Dict[str, str]], str, List[Tuple[str, str]], Any, List[Dict]]: history = history or [] chat_trace = chat_trace or [] message = user_message.strip() if not message: return ( tuples_to_messages(history), "", history, gr.update(value=user_message), chat_trace, ) strategy_key, strategy_name = resolve_embedding_selection(embedding_label) year_from, year_to = parse_year_range_input(year_range) model_id, temperature, extra_payload = resolve_reasoning_selection( preset_label, custom_model ) turn_settings = { "top_k": int(top_k), "year_range": [int(year_from), int(year_to)], "embedding_strategy": strategy_name, "embedding_key": strategy_key, "reasoning_profile": preset_label or "", "custom_model_input": (custom_model or "").strip(), "llm_model": model_id, "temperature": temperature, } hits = search_index( message, top_k=top_k, strategy_key=strategy_key, year_from=year_from, year_to=year_to, ) if not hits: reply = ( "No matching context was found for the selected year range and embedding strategy. " "Try widening the year filter or switching embeddings." ) updated_history = history + [(user_message, reply)] turn_record = { "timestamp": utc_now_iso_z(), "user_message": user_message, "assistant_reply": reply, "settings": turn_settings, "rag_hits": [], } updated_trace = chat_trace + [turn_record] return ( tuples_to_messages(updated_history), "", updated_history, gr.update(value=""), updated_trace, ) context_prompt = build_context_prompt(hits) system_message = { "role": "system", "content": ( "You are a research assistant. Use ONLY the provided context blocks when " "answering. Cite sources inline using square brackets matching the context " "indices (e.g., [1], [2]). If the information is unavailable, state that " "you do not know." ), } messages = [system_message] for prev_user, prev_bot in history: messages.append({"role": "user", "content": prev_user}) messages.append({"role": "assistant", "content": prev_bot}) augmented_user = ( f"Context:\n{context_prompt}\n\n" f"Current user message: {message}\n" "Respond with a concise answer followed by bullet points of key supporting facts. " f"Use citations [n] and preserve publication-year awareness (current retrieval embedding: {strategy_name})." ) messages.append({"role": "user", "content": augmented_user}) try: answer_text = CHAT_CLIENT.generate( messages, temperature=temperature, model_override=model_id, extra_payload=extra_payload, ) except Exception as exc: answer_text = ( "Failed to call the language model. " f"Verify the DEEPINFRA environment variable and model access. Details: {exc}" ) updated_history = history + [(user_message, answer_text)] references_text = format_references(hits) turn_record = { "timestamp": utc_now_iso_z(), "user_message": user_message, "assistant_reply": answer_text, "settings": turn_settings, "rag_hits": condense_hits_for_trace(hits), } updated_trace = chat_trace + [turn_record] return ( tuples_to_messages(updated_history), references_text, updated_history, gr.update(value=""), updated_trace, ) def reset_chat(): return [], "", [], gr.update(value=""), [] # --------------------------------------------------------------------------- # Gradio UI helpers # --------------------------------------------------------------------------- def registry_snapshot() -> Tuple[List[List], List[str]]: registry = load_registry_threadsafe() documents = registry.get("documents", {}) rows: List[List] = [] entries: List[str] = [] sorted_docs = sorted( documents.items(), key=lambda item: item[1].get("ingested_at", ""), reverse=True, ) for doc_id, doc in sorted_docs: year_value = doc.get("year") try: year_display = str(int(year_value)) except (TypeError, ValueError): year_display = "—" rows.append( [ "🗑️", doc_id, doc["title"], year_display, doc["pages"], doc.get("chunks", 0), ", ".join(map(str, doc.get("blank_pages", []))) or "—", doc.get("ingested_at", ""), ] ) entries.append(doc_id) return rows, entries def get_year_bounds() -> Tuple[int, int]: registry = load_registry_threadsafe() documents = registry.get("documents", {}) years: List[int] = [] for doc in documents.values(): try: years.append(int(doc.get("year"))) except (TypeError, ValueError): continue if years: return min(years), max(years) current_year = utc_year() return current_year - 20, current_year def year_slider_updates(selected_range: Any = None) -> Any: min_year, max_year = get_year_bounds() if min_year > max_year: min_year, max_year = max_year, min_year selected_from, selected_to = parse_year_range_input(selected_range) selected_from = max(min_year, min(max_year, int(selected_from))) selected_to = max(min_year, min(max_year, int(selected_to))) if selected_from > selected_to: selected_from, selected_to = selected_to, selected_from return gr.update( minimum=min_year, maximum=max_year, value=(selected_from, selected_to), ) def registry_outputs( selected_year_range: Any = None, ): table, _ = registry_snapshot() year_range_update = year_slider_updates(selected_year_range) return ( table, year_range_update, ) def handle_upload(files: List[str], selected_year_range: Any): if not files: table, year_range_update = registry_outputs( selected_year_range=selected_year_range ) return ( "Please upload at least one PDF.", table, year_range_update, ) messages = [] for file_path in files: try: result = ingest_pdf(file_path) messages.append(result["message"]) except Exception as exc: messages.append(f"Failed to index {Path(file_path).name}: {exc}") table, year_range_update = registry_outputs( selected_year_range=selected_year_range ) return "\n".join(messages), table, year_range_update def handle_delete(doc_id: Optional[str], selected_year_range: Any): doc_id = (doc_id or "").strip() if not doc_id: table, year_range_update = registry_outputs(selected_year_range=selected_year_range) return "Select a row in the document table to delete.", table, year_range_update try: outcome = delete_document(doc_id) except Exception as exc: outcome = { "status": "error", "doc_id": doc_id, "message": f"Delete failed for {doc_id}: {exc}", } table, year_range_update = registry_outputs(selected_year_range=selected_year_range) return outcome["message"], table, year_range_update def handle_refresh_registry(selected_year_range: Any): return registry_outputs(selected_year_range=selected_year_range) def _event_row_col(evt: gr.SelectData) -> Tuple[Optional[int], Optional[int]]: row_index: Optional[int] = None col_index: Optional[int] = None event_index = getattr(evt, "index", None) if isinstance(event_index, tuple): if len(event_index) >= 1: row_index = event_index[0] if len(event_index) >= 2: col_index = event_index[1] elif isinstance(event_index, list): if len(event_index) >= 1: row_index = event_index[0] if len(event_index) >= 2: col_index = event_index[1] elif isinstance(event_index, int): row_index = event_index return row_index, col_index def _row_from_table_data(table_data: Any, row_index: int) -> Optional[List[Any]]: if row_index < 0: return None # Gradio Dataframe can provide pandas-like or list-like structures. if hasattr(table_data, "iloc"): try: row = table_data.iloc[row_index] if hasattr(row, "tolist"): return list(row.tolist()) return list(row) except Exception: return None if isinstance(table_data, np.ndarray): try: row = table_data[row_index] return row.tolist() if hasattr(row, "tolist") else list(row) except Exception: return None if isinstance(table_data, list): if row_index >= len(table_data): return None row = table_data[row_index] if isinstance(row, (list, tuple)): return list(row) return [row] return None def handle_registry_row_select(table_data: Any, selected_year_range: Any, evt: gr.SelectData): try: row_index: Optional[int] = None col_index: Optional[int] = None row_index, col_index = _event_row_col(evt) selected_value = getattr(evt, "value", None) if row_index is None: return gr.update(), gr.update(), year_slider_updates(selected_year_range) if row_index < 0: return gr.update(), gr.update(), year_slider_updates(selected_year_range) row = _row_from_table_data(table_data, row_index) if not row or len(row) < 2: return gr.update(), gr.update(), year_slider_updates(selected_year_range) # First column acts like a delete action button. if col_index == 0 or selected_value == "🗑️": return handle_delete(doc_id=str(row[1]), selected_year_range=selected_year_range) return ( "Click the 🗑️ icon in the first column to delete this entry.", gr.update(), year_slider_updates(selected_year_range), ) except Exception as exc: return ( f"Delete action failed: {exc}", gr.update(), year_slider_updates(selected_year_range), ) def normalize_chat_history(history: Optional[List[Tuple[str, str]]]) -> List[Dict[str, str]]: normalized: List[Dict[str, str]] = [] for turn in history or []: if not isinstance(turn, (list, tuple)) or len(turn) < 2: continue normalized.append( { "user": str(turn[0]), "assistant": str(turn[1]), } ) return normalized def collect_rag_papers_from_trace(chat_trace: Optional[List[Dict]]) -> List[Dict]: papers: Dict[str, Dict] = {} for turn in chat_trace or []: if not isinstance(turn, dict): continue for hit in turn.get("rag_hits", []): if not isinstance(hit, dict): continue doc_id = str(hit.get("doc_id") or "").strip() doc_name = str(hit.get("doc_name") or "Document").strip() year_value = hit.get("year") key = doc_id or f"{doc_name}:{year_value}" if key in papers: continue papers[key] = { "doc_id": doc_id, "title": doc_name, "year": year_value, } return list(papers.values()) def feedback_snapshot() -> List[List]: store = load_feedback_registry_threadsafe() feedback_map = store.get("feedback", {}) if not isinstance(feedback_map, dict): return [] rows: List[List] = [] sorted_entries = sorted( feedback_map.items(), key=lambda item: item[1].get("created_at", "") if isinstance(item[1], dict) else "", reverse=True, ) for feedback_id, entry in sorted_entries: if not isinstance(entry, dict): continue settings = entry.get("settings", {}) if not isinstance(settings, dict): settings = {} chat_history = entry.get("chat_history", []) if not isinstance(chat_history, list): chat_history = [] rag_papers = entry.get("rag_papers", []) if not isinstance(rag_papers, list): rag_papers = [] rag_titles = [ str(paper.get("title") or "Document") for paper in rag_papers if isinstance(paper, dict) ] rag_display = ", ".join(rag_titles) if rag_titles else "—" year_range = settings.get("year_range") if isinstance(year_range, (list, tuple)) and len(year_range) >= 2: year_display = f"{year_range[0]}-{year_range[1]}" else: year_display = "—" status_value = str(entry.get("status") or "open") resolve_icon = "✅" if status_value != "resolved" else "☑️" rows.append( [ "🗑️", resolve_icon, str(feedback_id), status_value, str(entry.get("created_at") or ""), str(entry.get("resolved_at") or "—"), textwrap.shorten( str(entry.get("feedback_text") or ""), width=140, placeholder="…", ), len(chat_history), year_display, str(settings.get("embedding_strategy") or ""), str(settings.get("llm_model") or ""), textwrap.shorten(rag_display, width=120, placeholder="…"), ] ) return rows def feedback_entry_json(entry: Optional[Dict]) -> str: if not isinstance(entry, dict): return "" return json.dumps(entry, indent=2, ensure_ascii=True) def open_feedback_dialog(): return gr.update(visible=True), gr.update(value="") def close_feedback_dialog(): return gr.update(visible=False) def handle_submit_feedback( feedback_text: str, history: Optional[List[Tuple[str, str]]], chat_trace: Optional[List[Dict]], top_k: int, year_range: Any, embedding_label: Optional[str], preset_label: Optional[str], custom_model: Optional[str], ): message = (feedback_text or "").strip() if not message: feedback_rows = feedback_snapshot() warning = "Please enter feedback before submitting." return ( warning, gr.update(visible=True), gr.update(), feedback_rows, gr.update(), warning, ) strategy_key, strategy_name = resolve_embedding_selection(embedding_label) model_id, temperature, _ = resolve_reasoning_selection(preset_label, custom_model) year_from, year_to = parse_year_range_input(year_range) normalized_history = normalize_chat_history(history) rag_papers = collect_rag_papers_from_trace(chat_trace) llm_models_used = sorted( { str(turn.get("settings", {}).get("llm_model") or "").strip() for turn in (chat_trace or []) if isinstance(turn, dict) } | {model_id} ) llm_models_used = [model for model in llm_models_used if model] feedback_payload = { "feedback_id": "", "status": "open", "feedback_text": message, "created_at": utc_now_iso_z(), "resolved_at": "", "chat_history": normalized_history, "chat_trace": chat_trace or [], "settings": { "top_k": int(top_k), "year_range": [int(year_from), int(year_to)], "embedding_strategy": strategy_name, "embedding_key": strategy_key, "reasoning_profile": preset_label or "", "custom_model_input": (custom_model or "").strip(), "llm_model": model_id, "llm_models_used_in_chat": llm_models_used, "temperature": temperature, }, "rag_papers": rag_papers, } saved_entry = create_feedback_entry(feedback_payload) feedback_rows = feedback_snapshot() success = f"Feedback saved as {saved_entry['feedback_id']}." return ( success, gr.update(visible=False), gr.update(value=""), feedback_rows, gr.update(value=feedback_entry_json(saved_entry)), success, ) def handle_refresh_feedback(): return feedback_snapshot() def handle_feedback_row_select(table_data: Any, evt: gr.SelectData): try: row_index, col_index = _event_row_col(evt) selected_value = getattr(evt, "value", None) if row_index is None or row_index < 0: return "Select a feedback row.", gr.update(), gr.update() row = _row_from_table_data(table_data, row_index) if not row or len(row) < 3: return "Could not read selected feedback row.", gr.update(), gr.update() feedback_id = str(row[2]).strip() if not feedback_id: return "Feedback ID is missing in selected row.", gr.update(), gr.update() if col_index == 0 or selected_value == "🗑️": result = delete_feedback_entry(feedback_id) return result["message"], feedback_snapshot(), gr.update(value="") if col_index == 1 or selected_value in {"✅", "☑️"}: result = mark_feedback_resolved(feedback_id) entry = get_feedback_entry(feedback_id) detail_json = feedback_entry_json(entry) return result["message"], feedback_snapshot(), gr.update(value=detail_json) entry = get_feedback_entry(feedback_id) if not entry: return f"Feedback {feedback_id} not found.", feedback_snapshot(), gr.update(value="") return ( f"Selected feedback {feedback_id}.", gr.update(), gr.update(value=feedback_entry_json(entry)), ) except Exception as exc: return f"Feedback action failed: {exc}", gr.update(), gr.update() def build_interface() -> gr.Blocks: ensure_strategy_indexes_initialized() initial_table, _ = registry_snapshot() initial_feedback_table = feedback_snapshot() initial_year_min, initial_year_max = get_year_bounds() with gr.Blocks( title="Document RAG", css=""" .feedback-dialog-panel { border: 1px solid #d1d5db; border-radius: 12px; padding: 12px; margin-top: 10px; background: #f9fafb; } """, ) as demo: gr.Markdown( """ # Document RAG Use the tabs: - **Upload PDFs**: index/delete RAG documents. - **Chat**: query the indexed documents with year-range and embedding controls. - **Feedback**: review, resolve, and delete feedback submissions. - **Config**: environment variable reference. """ ) with gr.Tab("Upload PDFs"): with gr.Row(): file_input = gr.File( label="PDF files", file_types=[".pdf"], file_count="multiple", type="filepath", ) upload_button = gr.Button("Process and Index") upload_status = gr.Textbox( label="Status", lines=4, interactive=False, show_label=True ) registry_table = gr.Dataframe( headers=[ "Delete", "Doc ID", "Title", "Year", "Pages", "Chunks", "OCR pages", "Indexed at", ], datatype=["str", "str", "str", "str", "number", "number", "str", "str"], interactive=False, value=initial_table, ) with gr.Row(): refresh_button = gr.Button("Refresh List") with gr.Tab("Chat"): reasoning_options = list(REASONING_PRESETS.keys()) default_preset = reasoning_options[0] if reasoning_options else None with gr.Row(): chatbot = gr.Chatbot( label="Conversation", height=420, show_label=True, sanitize_html=False, type="messages", ) references_md = gr.Markdown(label="References (last reply)") with gr.Row(): chat_input = gr.Textbox( label="Message", placeholder="Ask about the indexed PDFs…", lines=2, ) send_button = gr.Button("Send", variant="primary") clear_button = gr.Button("Clear", variant="secondary") feedback_button = gr.Button("feedback", variant="secondary") with gr.Row(): top_k_slider = gr.Slider( minimum=1, maximum=20, value=10, step=1, label="Top-k documents", ) if RangeSlider is not None: year_range_slider = RangeSlider( minimum=initial_year_min, maximum=initial_year_max, value=(initial_year_min, initial_year_max), step=1, label="Year range", interactive=True, ) else: # Fallback to built-in slider when the range-slider plugin is unavailable. year_range_slider = gr.Slider( minimum=initial_year_min, maximum=initial_year_max, value=(initial_year_min, initial_year_max), step=1, label="Year range", interactive=True, ) with gr.Row(): embedding_dropdown = gr.Dropdown( label="Retrieval embedding strategy", choices=EMBEDDING_OPTIONS, value=DEFAULT_EMBEDDING_LABEL, allow_custom_value=False, ) reasoning_dropdown = gr.Dropdown( label="Reasoning profile", choices=reasoning_options, value=default_preset, allow_custom_value=False, ) custom_model_box = gr.Textbox( label="Custom model ID (optional)", placeholder=f"Overrides preset (default: {DEFAULT_CHAT_MODEL})", ) chat_feedback_status = gr.Textbox( label="Feedback status", lines=1, interactive=False, show_label=True, value="", ) with gr.Group(visible=False, elem_classes=["feedback-dialog-panel"]) as feedback_dialog: gr.Markdown("### Submit feedback") feedback_input = gr.Textbox( label="What should be fixed?", lines=4, placeholder="Describe what happened and what should change.", ) with gr.Row(): submit_feedback_button = gr.Button("Submit feedback", variant="primary") cancel_feedback_button = gr.Button("Cancel", variant="secondary") history_state = gr.State([]) chat_trace_state = gr.State([]) with gr.Tab("Feedback"): feedback_tab_status = gr.Textbox( label="Status", lines=2, interactive=False, show_label=True, value="", ) feedback_table = gr.Dataframe( headers=[ "Delete", "Resolve", "Feedback ID", "Status", "Created at", "Resolved at", "Feedback", "Turns", "Year range", "Embedding", "LLM", "RAG papers", ], datatype=[ "str", "str", "str", "str", "str", "str", "str", "number", "str", "str", "str", "str", ], interactive=False, value=initial_feedback_table, ) with gr.Row(): feedback_refresh_button = gr.Button("Refresh Feedback") gr.Markdown( "Click `✅` in the Resolve column to mark an item as resolved, or `🗑️` to delete it." ) feedback_detail_box = gr.Textbox( label="Selected feedback payload (JSON)", lines=18, interactive=False, ) with gr.Tab("Config"): gr.Markdown( """ ## Runtime Configuration - `APP_URL_TOKEN` (optional): URL token for app access (used as `?token=...` by default). - `APP_ENABLE_TOKEN_GATE_ON_SPACES` (optional): defaults to `0`; set to `1` only if you explicitly want query-token gating on Spaces. - `APP_ALLOW_LOOPBACK_BYPASS` (optional): defaults to `1`; keeps localhost startup/health checks working. - `DEEPINFRA`: required for OCR and answer generation. - `DEEPINFRA_OCR_MODEL` (optional): defaults to `deepseek-ai/deepseek-ocr`. - `DEEPINFRA_RAG_MODEL` (optional): defaults to `meta-llama/Meta-Llama-3.1-8B-Instruct`. - `DEEPINFRA_FAST_MODEL`, `DEEPINFRA_BALANCED_MODEL`, `DEEPINFRA_REASONING_MODEL` (optional): override preset models. - `EMBED_MODEL_FAST`, `EMBED_MODEL_QUALITY`, `EMBED_MODEL_QA` (optional): override retrieval embedding models. """ ) if RangeSlider is None: gr.Markdown( "Year range uses fallback slider because `gradio_rangeslider` is not installed." ) send_inputs = [ chat_input, history_state, top_k_slider, year_range_slider, embedding_dropdown, reasoning_dropdown, custom_model_box, chat_trace_state, ] send_outputs = [chatbot, references_md, history_state, chat_input, chat_trace_state] registry_outputs_all = [upload_status, registry_table, year_range_slider] registry_refresh_outputs = [registry_table, year_range_slider] feedback_submit_outputs = [ chat_feedback_status, feedback_dialog, feedback_input, feedback_table, feedback_detail_box, feedback_tab_status, ] feedback_row_outputs = [ feedback_tab_status, feedback_table, feedback_detail_box, ] upload_button.click( fn=handle_upload, inputs=[file_input, year_range_slider], outputs=registry_outputs_all, ) refresh_button.click( fn=handle_refresh_registry, inputs=[year_range_slider], outputs=registry_refresh_outputs, ) registry_table.select( fn=handle_registry_row_select, inputs=[registry_table, year_range_slider], outputs=registry_outputs_all, ) send_button.click( fn=chat_with_rag, inputs=send_inputs, outputs=send_outputs, ) chat_input.submit( fn=chat_with_rag, inputs=send_inputs, outputs=send_outputs, ) clear_button.click( fn=reset_chat, outputs=send_outputs, queue=False, ) feedback_button.click( fn=open_feedback_dialog, outputs=[feedback_dialog, chat_feedback_status], queue=False, ) cancel_feedback_button.click( fn=close_feedback_dialog, outputs=[feedback_dialog], queue=False, ) submit_feedback_button.click( fn=handle_submit_feedback, inputs=[ feedback_input, history_state, chat_trace_state, top_k_slider, year_range_slider, embedding_dropdown, reasoning_dropdown, custom_model_box, ], outputs=feedback_submit_outputs, ) feedback_refresh_button.click( fn=handle_refresh_feedback, outputs=[feedback_table], ) feedback_table.select( fn=handle_feedback_row_select, inputs=[feedback_table], outputs=feedback_row_outputs, ) demo.load( fn=handle_refresh_registry, inputs=[year_range_slider], outputs=registry_refresh_outputs, queue=False, ) demo.load( fn=handle_refresh_feedback, outputs=[feedback_table], queue=False, ) return demo if __name__ == "__main__": app = build_interface() launch_kwargs: Dict[str, Any] = {"ssr_mode": False} if URL_ACCESS_TOKEN and not TOKEN_GATE_ENABLED: print( "APP_URL_TOKEN is set but token gating is disabled on Spaces by default. " "Set APP_ENABLE_TOKEN_GATE_ON_SPACES=1 to enable it." ) if TOKEN_GATE_ENABLED: launch_kwargs["app_kwargs"] = { "middleware": [Middleware(QueryTokenGateMiddleware)] } app.launch(**launch_kwargs)