Blame for app/query.py

"""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
@dataclass
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