Spaces:
Running on Zero
Running on Zero
| """Shared CLI/UI answering with persistent answer cache and bounded API usage.""" | |
| import hashlib | |
| import json | |
| import sqlite3 | |
| from collections import deque | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from threading import RLock | |
| from time import monotonic | |
| from visual_doc_rag.data import write_json | |
| from visual_doc_rag.generate import ( | |
| BASE_URL, | |
| PROMPT_VERSION, | |
| GenerationError, | |
| Generator, | |
| Settings, | |
| validate_answer, | |
| ) | |
| from visual_doc_rag.retrieve import PageRetriever | |
| def answer_cache_key(query, method, source, model, pages): | |
| fingerprint = { | |
| "query": query, | |
| "method": method, | |
| "dataset": source, | |
| "model": model, | |
| "endpoint": BASE_URL, | |
| "prompt": PROMPT_VERSION, | |
| "image_processing": "JPEG quality=90, max_side=1600", | |
| "images": [ | |
| [p["corpus_id"], hashlib.sha256(Path(p["image_path"]).read_bytes()).hexdigest()] | |
| for p in pages | |
| ], | |
| } | |
| return hashlib.sha256(json.dumps(fingerprint, sort_keys=True).encode()).hexdigest() | |
| class RequestBudget: | |
| """Count every outbound attempt, including retries/errors, across local processes.""" | |
| def __init__(self, root: Path, limit: int): | |
| root.mkdir(parents=True, exist_ok=True) | |
| self.path = root / "usage.sqlite3" | |
| self.limit = limit | |
| with sqlite3.connect(self.path) as db: | |
| db.execute("CREATE TABLE IF NOT EXISTS usage (day TEXT PRIMARY KEY, requests INTEGER)") | |
| def reserve(self): | |
| day = datetime.now(UTC).date().isoformat() | |
| with sqlite3.connect(self.path) as db: | |
| db.execute("BEGIN IMMEDIATE") | |
| row = db.execute("SELECT requests FROM usage WHERE day = ?", (day,)).fetchone() | |
| if row and row[0] >= self.limit: | |
| raise GenerationError( | |
| "Local daily API request cap reached. Try again tomorrow (UTC)." | |
| ) | |
| db.execute( | |
| "INSERT INTO usage VALUES (?, 1) ON CONFLICT(day) " | |
| "DO UPDATE SET requests = requests + 1", | |
| (day,), | |
| ) | |
| class AnswerService: | |
| def __init__( | |
| self, subset="hr", data_dir=Path("data"), settings=None, retriever=None, generator=None | |
| ): | |
| self.settings = settings or Settings.from_env() | |
| self.subset, self.data_dir = subset, data_dir | |
| self.retriever = retriever | |
| self.generator = generator or Generator(self.settings) | |
| self.cache_root = data_dir / subset / "answers" | |
| self.budget = RequestBudget(data_dir / "runtime", self.settings.daily_requests) | |
| self.sessions = {} | |
| self.lock = RLock() | |
| def ask(self, query: str, method="hybrid", session_id="cli", *, cache_only=False) -> dict: | |
| query = " ".join(query.split()) | |
| if not query or len(query) > 2000: | |
| raise GenerationError("Enter a question of 1–2,000 characters.") | |
| with self.lock: | |
| # Cache normalized whitespace, preserving case and punctuation. | |
| if self.retriever is None: | |
| self.retriever = PageRetriever(self.subset, self.data_dir) | |
| pages = self.retriever.search(query, method, top_k=3) | |
| key = answer_cache_key(query, method, self.retriever.source, self.settings.model, pages) | |
| cache = self.cache_root / f"{key}.json" | |
| if cache.exists(): | |
| try: | |
| answer = validate_answer(cache.read_text(encoding="utf-8"), pages) | |
| cached = True | |
| except (ValueError, OSError): | |
| cached = False | |
| else: | |
| cached = False | |
| if not cached: | |
| if cache_only: | |
| raise GenerationError( | |
| "No valid cached answer for this question and configuration." | |
| ) | |
| now = monotonic() | |
| self.sessions = { | |
| key: deque(t for t in times if now - t < 600) | |
| for key, times in self.sessions.items() | |
| if times and now - times[-1] < 600 | |
| } | |
| times = self.sessions.setdefault(session_id, deque()) | |
| if len(times) >= 5: | |
| raise GenerationError("Session limit reached: 5 new questions per 10 minutes.") | |
| times.append(now) | |
| answer = self.generator.generate(query, pages, self.budget.reserve) | |
| write_json(cache, answer.model_dump()) | |
| return { | |
| "query": query, | |
| "method": method, | |
| "model": self.settings.model, | |
| **answer.model_dump(), | |
| "cached": cached, | |
| "pages": pages, | |
| } | |