Spaces:
Runtime error
Runtime error
| """Understand the question BEFORE searching: page numbers, which document, a clean search query.""" | |
| import json | |
| import logging | |
| import re | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from . import llm | |
| from .config import settings | |
| log = logging.getLogger("rag.query") | |
| # ------------------------------------------------------------------ page numbers | |
| _NUM = r"\d{1,4}" | |
| _SEP = r"(?:-|\u2013|\u2014|to|se|,|&|and|aur)" | |
| _LIST = rf"{_NUM}(?:\s*{_SEP}\s*{_NUM})*" | |
| # "page 20", "pages 20-25", "page no. 20", "pg 5, 6 and 7" | |
| _PAGE_AFTER = re.compile( | |
| rf"(?:\bpages?|\bpgs?\b|\bpp\.)\s*(?:no\.?|number|num|nbr|#|:)?\s*({_LIST})", re.I | |
| ) | |
| # "20th page", "20 no ke page", "20 se 25 pages" | |
| _PAGE_BEFORE = re.compile( | |
| rf"\b({_NUM}(?:\s*(?:-|\u2013|\u2014|to|se)\s*{_NUM})?)(?:st|nd|rd|th)?\s*" | |
| r"(?:no\.?|number|nambar|num)?\s*(?:ke|ka|ki|wale|wala|vale|vala|wali|vali)?\s*" | |
| r"(?:pages?|pg|panne|panna|pej)\b", | |
| re.I, | |
| ) | |
| def _parse_expr(expr: str) -> set[int]: | |
| expr = re.sub(r"\b(?:to|se)\b|\u2013|\u2014", "-", expr, flags=re.I) | |
| out: set[int] = set() | |
| for part in re.split(r"\s*(?:,|&|\band\b|\baur\b)\s*", expr, flags=re.I): | |
| part = part.strip() | |
| m = re.fullmatch(r"(\d+)\s*-\s*(\d+)", part) | |
| if m: | |
| a, b = int(m[1]), int(m[2]) | |
| if a <= b and b - a <= 60: | |
| out.update(range(a, b + 1)) | |
| else: | |
| out.update([a, b]) | |
| elif part.isdigit(): | |
| out.add(int(part)) | |
| return out | |
| def extract_pages(question: str) -> list[int]: | |
| pages: set[int] = set() | |
| for rx in (_PAGE_AFTER, _PAGE_BEFORE): | |
| for m in rx.finditer(question): | |
| pages |= _parse_expr(m.group(1)) | |
| return sorted(p for p in pages if p >= 1)[:50] | |
| # --------------------------------------------------------------- document names | |
| def _norm(s: str) -> str: | |
| return re.sub(r"[\W_]+", " ", s.lower()).strip() | |
| def match_docs(question: str, docs: list[dict]) -> list[str]: | |
| """Documents whose file name is mentioned in the question.""" | |
| q = f" {_norm(question)} " | |
| hits = [] | |
| for d in docs: | |
| stem = _norm(Path(d["filename"]).stem) | |
| if len(stem) >= 3 and f" {stem} " in q: | |
| hits.append(d["doc_id"]) | |
| return hits | |
| # ---------------------------------------------------------------- LLM rewriting | |
| REWRITE_SYSTEM = """You prepare a user's question for a document search engine. You do NOT answer it. | |
| The user may write English, Hindi or Hinglish (Hindi in Roman letters), and may have typos. | |
| Return ONLY a JSON object, no other text: | |
| {"search_query": "...", "pages": [], "document": null, "intent": "search"} | |
| - search_query: a standalone, keyword-rich ENGLISH search query. Fix typos, translate Hinglish, | |
| and resolve words like "it" or "that" using recent_chat. Do not include page numbers in it. | |
| - pages: list of page numbers the user refers to (e.g. "20 no ke page" -> [20]); [] if none. | |
| - document: one name copied EXACTLY from available_documents if the user clearly refers to a | |
| specific document, otherwise null. | |
| - intent: "page_lookup" if the user wants to know what is written on / contained in specific pages, | |
| otherwise "search".""" | |
| def _extract_json(raw: str) -> dict: | |
| start, end = raw.find("{"), raw.rfind("}") | |
| if start == -1 or end <= start: | |
| raise ValueError("no JSON in model reply") | |
| return json.loads(raw[start : end + 1]) | |
| def _rewrite_with_llm(question: str, docs: list[dict], history: list[dict]) -> dict: | |
| recent = [{"role": h["role"], "text": h["content"][:300]} for h in history[-4:]] | |
| payload = json.dumps( | |
| {"available_documents": [d["filename"] for d in docs][:30], "recent_chat": recent, "question": question}, | |
| ensure_ascii=False, | |
| ) | |
| raw = llm.chat( | |
| [{"role": "system", "content": REWRITE_SYSTEM}, {"role": "user", "content": payload}], | |
| max_tokens=800, | |
| temperature=0, | |
| ) | |
| return _extract_json(raw) | |
| # ------------------------------------------------------------------------- plan | |
| class QueryPlan: | |
| search_query: str | |
| intent: str = "search" # "search" | "page_lookup" | |
| pages: list[int] = field(default_factory=list) | |
| doc_ids: list[str] = field(default_factory=list) # documents named in the question | |
| used_llm: bool = False | |
| def analyze(question: str, docs: list[dict], history: list[dict] | None = None) -> QueryPlan: | |
| history = history or [] | |
| regex_pages = extract_pages(question) | |
| doc_ids = match_docs(question, docs) | |
| plan = QueryPlan(search_query=question, pages=regex_pages, doc_ids=doc_ids) | |
| plan.intent = "page_lookup" if regex_pages else "search" | |
| if not settings.use_query_rewrite: | |
| return plan | |
| try: | |
| data = _rewrite_with_llm(question, docs, history) | |
| except Exception as e: # rate limit, bad JSON, no key... fall back to the plain question | |
| log.warning("Query rewrite skipped: %s", e) | |
| return plan | |
| plan.used_llm = True | |
| sq = data.get("search_query") | |
| if isinstance(sq, str) and sq.strip(): | |
| plan.search_query = sq.strip() | |
| llm_pages = [int(p) for p in (data.get("pages") or []) if str(p).isdigit() and int(p) >= 1][:50] | |
| plan.pages = sorted(set(llm_pages)) or regex_pages | |
| if plan.pages: | |
| plan.intent = "page_lookup" if data.get("intent") != "search" else "search" | |
| else: | |
| plan.intent = "search" | |
| if not plan.doc_ids and isinstance(data.get("document"), str): | |
| wanted = data["document"].strip().lower() | |
| for d in docs: | |
| if d["filename"].lower() == wanted: | |
| plan.doc_ids = [d["doc_id"]] | |
| break | |
| return plan | |