WiseGrade/backend/wisegrade/rag.py

"""Indexing and hybrid retrieval with page-type aware candidate filtering.

Chinese keyword recall does not use PostgreSQL ``to_tsvector('simple', …)``. That
configuration has no Chinese segmenter, so a whole run of Chinese characters becomes a
single token; measured on the course reports, the query ``文献调研`` does not match the
page that contains ``第2章文献调研与技术现状``. Recall therefore uses character bigrams
plus whole ASCII/number tokens, which needs no extension and behaves the same on
PostgreSQL and SQLite.

Semantic and keyword scores are fused by rank (reciprocal rank fusion), never by adding
raw numbers: a cosine similarity and a bigram overlap are not on a common scale.
"""
from functools import lru_cache
from pathlib import Path
import re

from sqlalchemy import case, cast, func, literal, or_, select

from .config import get_settings
from .models import KnowledgeChunk, KnowledgeDocument, KnowledgePage, SubmissionFile
from .knowledge_pages import page_text, relative_image_path
from .pages import (BLOCK_BODY, BLOCK_TOC, is_toc_entry,
                    normalize_heading, searchable_pages)
from .evidence import heading_only_text

INDEX_VERSION = 5

#: Fusion constant for reciprocal rank fusion; the usual value from the IR literature.
RRF_K = 60
#: Candidate pools are gathered independently so both retrievers get a chance.
SEMANTIC_POOL = 60
KEYWORD_POOL = 60
#: A chunk must be at least this long to count as substantive content evidence.
MIN_CHUNK_CHARS = 60
#: At most this many chunks from one page in the first retrieval pass.
DEFAULT_PAGE_LIMIT = 2
#: Soft cap on chunks from one section once a reliable section path exists.
DEFAULT_SECTION_CAP = 3
#: Maximum chunks handed to the model for one criterion, across all targets.
SCORE_EVIDENCE_LIMIT = 14

_ASCII_TOKEN = re.compile(r"[A-Za-z][A-Za-z0-9_+\-.]*|\d+(?:\.\d+)*")
_CJK_RUN = re.compile(r"[㐀-䶿一-鿿]+")


def query_terms(query: str) -> tuple[list[str], list[str]]:
    """Split a query into (ascii tokens lowercased, CJK character bigrams).

    Single-character CJK runs are dropped: they carry too little signal and would make
    every chunk with a common character look relevant.
    """
    value = str(query or "")
    ascii_tokens = [token.lower() for token in _ASCII_TOKEN.findall(value)]
    bigrams: list[str] = []
    seen = set()
    for run in _CJK_RUN.findall(value):
        for index in range(len(run) - 1):
            bigram = run[index:index + 2]
            if bigram not in seen:
                seen.add(bigram)
                bigrams.append(bigram)
    return sorted(set(ascii_tokens)), bigrams


def split_pages(pages: list[dict], *, size=650, overlap=100):
    """Split pages into searchable chunks carrying page type and section path."""
    if size <= 0 or overlap < 0 or overlap >= size:
        raise ValueError("切块大小必须为正且重叠小于块大小")
    for page in pages:
        page_type = page.get("page_type") or "body"
        chapter = page.get("chapter", "")
        section = page.get("section", "")
        blocks = page.get("blocks")
        if not blocks:
            blocks = [{"kind": BLOCK_BODY, "text": str(page.get("text") or "").strip()}]
            blocks.extend({"kind": BLOCK_BODY, "text": "表格: " + str(table.get("rows", table))}
                          for table in page.get("tables", []))
        for block in blocks:
            source_type = block.get("page_type", page_type)
            block_chapter = block.get("chapter", chapter)
            block_section = block.get("section", section)
            if block["kind"] == BLOCK_TOC:
                # A contents page is navigation, so its text is not body content.
                yield {"page": page["page_number"], "page_type": "toc", "block_type": BLOCK_TOC,
                       "chapter": chapter, "section": section, "text": block["text"]}
                continue
            text_block = block["text"].strip()
            if not text_block:
                continue
            for table in ([block["tables"]] if block.get("tables") else []):
                for table_item in table:
                    text_block += "\n表格: " + str(table_item.get("rows", table_item))
            # Table rows of a page are appended once, on the page's first body chunk.
            start = 0
            while start < len(text_block):
                part = text_block[start:start + size]
                unsafe = any(c in part for c in ("\ufffd", "\ufffe", "\uffff")) or heading_only_text(part)
                safe_type = "unknown" if unsafe and source_type in ("body", "appendix") else source_type
                yield {"page": page["page_number"], "page_type": safe_type, "block_type": BLOCK_BODY,
                       "chapter": block_chapter, "section": block_section, "text": part}
                if start + size >= len(text_block):
                    break
                start += size - overlap


def build_chunks(pages: list[dict]) -> list[dict]:
    """Chunks for one document, classifying pages first when they are not yet typed."""
    prepared = pages if pages and "page_type" in (pages[0] or {}) else searchable_pages(pages)
    return list(split_pages(prepared))


@lru_cache(maxsize=2)
def local_encoder(model_path: str):
    path = Path(model_path)
    if not path.is_dir() or not any(path.iterdir()):
        raise RuntimeError(f"离线嵌入模型不存在: {path}；请挂载本地模型包")
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer(str(path), local_files_only=True, trust_remote_code=False)


def embed(texts: list[str], *, model_path: str | None = None):
    if not texts:
        return []
    model = local_encoder(model_path or str(get_settings().embedding_model_path))
    vectors = model.encode(texts, normalize_embeddings=True)
    output = []
    for vector in vectors:
        values = vector.tolist()
        if len(values) > 1024:
            raise RuntimeError("嵌入维度超过pgvector列容量1024")
        output.append(values + [0.0] * (1024 - len(values)))
    return output


def ingest(db, *, course_version_id: int, project_id: int | None, submission_file_id: int | None,
           title: str, category: str, source_hash: str, pages: list[dict], embedder=embed):
    history = db.scalars(select(KnowledgeDocument).where(KnowledgeDocument.course_version_id == course_version_id,
                    KnowledgeDocument.project_id == project_id, KnowledgeDocument.submission_file_id == submission_file_id,
                    KnowledgeDocument.title == title).order_by(KnowledgeDocument.id.desc())).all()
    existing = next((doc for doc in history if doc.source_hash == source_hash and doc.status == "ready"
                     and doc.index_version == INDEX_VERSION and doc.document_type == category), None)
    if existing:
        return existing
    prepared = searchable_pages(pages)
    chunks = list(split_pages(prepared))
    vectors = embedder([chunk["text"] for chunk in chunks])
    if len(vectors) != len(chunks):
        raise RuntimeError("嵌入向量数量与知识片段数量不一致")
    # Pages and chunk IDs are immutable citation sources. Even an unchanged file
    # needs a new document revision when its index or material category changes.
    revision = (db.scalar(select(func.max(KnowledgeDocument.revision)).where(
        KnowledgeDocument.course_version_id == course_version_id, KnowledgeDocument.project_id == project_id,
        KnowledgeDocument.submission_file_id == submission_file_id, KnowledgeDocument.source_hash == source_hash)) or 0) + 1
    doc = KnowledgeDocument(course_version_id=course_version_id, project_id=project_id, submission_file_id=submission_file_id,
                            title=title, document_type=category, source_hash=source_hash, revision=revision)
    db.add(doc)
    doc.index_version, doc.status = INDEX_VERSION, "indexing"
    db.flush()
    for source, prepared_page in zip(pages, prepared):
        classification = prepared_page["classification"]
        db.add(KnowledgePage(
            document_id=doc.id, page=source["page_number"], text=page_text(source),
            image_path=relative_image_path(source.get("rendered_image_path")),
            page_type=classification["page_type"], mixed=bool(classification["mixed"]),
            section=classification["section"], chapter=classification["chapter"],
            confidence=classification["confidence"],
            classification={**classification["signals"], "classifier_version": classification["classifier_version"],
                            "blocks": [{"kind": block["kind"], "chars": len(block["text"]),
                                        "substantive": block.get("substantive", False)}
                                       for block in classification["blocks"]]}))
    for chunk, vector in zip(chunks, vectors):
        db.add(KnowledgeChunk(document_id=doc.id, course_version_id=course_version_id, project_id=project_id,
                              document_type=category, page=chunk["page"], section=chunk["section"],
                              chapter=chunk["chapter"], page_type=chunk["page_type"],
                              block_type=chunk["block_type"], text=chunk["text"], embedding=vector))
    for old in history:
        old.status = "superseded"
    doc.status = "ready"
    db.commit()
    return doc


def _scope_filter(project_id: int | None, include_public: bool, public_only: bool):
    if public_only:
        return KnowledgeChunk.project_id.is_(None)
    if include_public:
        return or_(KnowledgeChunk.project_id == project_id, KnowledgeChunk.project_id.is_(None))
    return KnowledgeChunk.project_id == project_id


def current_document_ids():
    selected_file = select(SubmissionFile.id).where(
        SubmissionFile.id == KnowledgeDocument.submission_file_id,
        SubmissionFile.recognition_status == "recognized", SubmissionFile.parse_status == "ready",
        SubmissionFile.sha256 == KnowledgeDocument.source_hash).exists()
    return select(KnowledgeDocument.id).where(KnowledgeDocument.status == "ready",
        KnowledgeDocument.index_version == INDEX_VERSION,
        or_(KnowledgeDocument.submission_file_id.is_(None), selected_file))


def _base_filters(*, course_version_id: int, project_id: int | None, category, categories,
                  include_public: bool, public_only: bool, allowed_page_types, exclude_document_ids, document_id=None):
    filters = [KnowledgeChunk.course_version_id == course_version_id,
               _scope_filter(project_id, include_public, public_only),
               KnowledgeChunk.block_type != "legacy",
               KnowledgeChunk.document_id.in_(current_document_ids())]
    if document_id is not None:
        filters.append(KnowledgeChunk.document_id == document_id)
    if category:
        filters.append(KnowledgeChunk.document_type == category)
    if categories:
        filters.append(KnowledgeChunk.document_type.in_(list(categories)))
    if allowed_page_types is not None:
        filters.append(KnowledgeChunk.page_type.in_(list(allowed_page_types)))
    if exclude_document_ids:
        filters.append(KnowledgeChunk.document_id.not_in(list(exclude_document_ids)))
    return filters


def _keyword_query(db, filters, query: str, limit: int):
    """Keyword recall ranked by bigram/token overlap over the already-filtered rows.

    PostgreSQL's ``to_tsvector('simple', …)`` cannot segment Chinese, so it is not used
    for recall. The overlap is computed with ``strpos``/``substr`` so the same expression
    works on every dialect SQLAlchemy targets here.
    """
    ascii_tokens, bigrams = query_terms(query)
    if not ascii_tokens and not bigrams:
        return []
    if db.bind.dialect.name != "postgresql":
        rows = db.execute(select(KnowledgeChunk.id, KnowledgeChunk.text).where(*filters)).all()
        ranked = [(row[0], keyword_overlap(row[1], ascii_tokens, bigrams)) for row in rows]
        ranked.sort(key=lambda row: (-row[1], row[0]))
        return [chunk_id for chunk_id, score in ranked[:limit] if score > 0]
    total = len(ascii_tokens) + len(bigrams)
    score = None
    for token in ascii_tokens:
        pattern = r"(^|[^a-z0-9_+.-])" + re.escape(token) + r"($|[^a-z0-9_+.-])"
        hit = case((func.lower(KnowledgeChunk.text).op("~")(pattern), 1.0), else_=0.0)
        score = hit if score is None else score + hit
    if bigrams:
        # Count each query term once, matching keyword_overlap and avoiding a bonus
        # for repeated words on an otherwise irrelevant page.
        for bigram in bigrams:
            hit = case((func.strpos(KnowledgeChunk.text, literal(bigram)) > 0, 1.0), else_=0.0)
            score = hit if score is None else score + hit
    score = (score / float(total)).label("keyword")
    rows = db.execute(select(KnowledgeChunk.id, score).where(*filters)
                      .order_by(score.desc(), KnowledgeChunk.id.asc()).limit(limit)).all()
    return [row[0] for row in rows if float(row[1] or 0) > 0]


def keyword_overlap(text_value: str, ascii_tokens: list[str], bigrams: list[str]) -> float:
    """Bigram/token overlap used by the non-PostgreSQL path and by tests."""
    value = str(text_value or "")
    lower = value.lower()
    tokens = set(_ASCII_TOKEN.findall(lower))
    hits = sum(1 for token in ascii_tokens if token in tokens)
    total = len(ascii_tokens)
    hits += sum(1 for bigram in bigrams if bigram in value)
    total += len(bigrams)
    return hits / total if total else 0.0


def _semantic_query(db, filters, vector, limit: int):
    if db.bind.dialect.name == "postgresql":
        from pgvector.sqlalchemy import Vector
        distance = cast(KnowledgeChunk.embedding, Vector(1024)).cosine_distance(vector)
        rows = db.execute(select(KnowledgeChunk.id).where(*filters)
                          .order_by(distance.asc(), KnowledgeChunk.id.asc()).limit(limit)).all()
        return [row[0] for row in rows]
    rows = db.execute(select(KnowledgeChunk.id, KnowledgeChunk.embedding).where(*filters)).all()
    scored = [(chunk_id, sum(a * b for a, b in zip(embedding or [], vector)))
              for chunk_id, embedding in rows]
    scored.sort(key=lambda pair: (-pair[1], pair[0]))
    return [chunk_id for chunk_id, _ in scored[:limit]]


def _fuse(db, ids: list[int], semantic_ids: list[int], keyword_ids: list[int]):
    """Reciprocal rank fusion over the two ranked lists, keeping discovery order stable."""
    scores: dict[int, float] = {chunk_id: 0.0 for chunk_id in ids}
    for ranked in (semantic_ids, keyword_ids):
        for rank, chunk_id in enumerate(ranked):
            if chunk_id in scores:
                scores[chunk_id] += 1.0 / (RRF_K + rank + 1)
    return scores


def _row_dict(chunk):
    return {"chunk_id": chunk.id, "document_id": chunk.document_id, "project_id": chunk.project_id,
            "document_type": chunk.document_type, "page": chunk.page, "section": chunk.section,
            "chapter": chunk.chapter, "page_type": chunk.page_type, "block_type": chunk.block_type,
            "text": chunk.text}


def normalize_hit(hit: dict) -> dict:
    """Fill the fields the rest of the pipeline reads, so one shape flows through.

    An alternate retriever or a test double may return only the essentials. Filling the
    defaults here keeps every consumer from re-deriving them; the authoritative page type
    is still re-read from storage when a score is produced.
    """
    return {"chunk_id": hit.get("chunk_id"), "document_id": hit.get("document_id"),
            "project_id": hit.get("project_id"), "document_type": hit.get("document_type"),
            "page": hit.get("page"), "text": hit.get("text") or "",
            "section": hit.get("section") or "", "chapter": hit.get("chapter") or "",
            "page_type": hit.get("page_type") or "unknown", "block_type": hit.get("block_type") or "legacy",
            "rank": float(hit.get("rank") or 0.0)}


def retrieve(db, *, course_version_id: int, project_id: int, query: str, category: str | None = None,
             categories=None, limit=10, embedder=embed, include_public=True, public_only=False,
             allowed_page_types=None, exclude_document_ids=(), substantive_only=False,
             semantic_pool: int | None = None, keyword_pool: int | None = None, document_id=None):
    """Filtered hybrid retrieval, then rank fusion over the two independent recalls.

    Page filtering happens inside the query, so it always precedes the candidate limit:
    a contents page can never displace a body page simply by scoring well.
    """
    filters = _base_filters(course_version_id=course_version_id, project_id=project_id, category=category,
                            categories=categories, include_public=include_public, public_only=public_only,
                            allowed_page_types=allowed_page_types, exclude_document_ids=exclude_document_ids,
                            document_id=document_id)
    if substantive_only:
        filters.append(KnowledgeChunk.block_type != BLOCK_TOC)
    vector = embedder([query])[0]
    semantic_ids = _semantic_query(db, filters, vector, semantic_pool or SEMANTIC_POOL)
    keyword_ids = _keyword_query(db, filters, query, keyword_pool or KEYWORD_POOL)
    ids = list(dict.fromkeys(semantic_ids + keyword_ids))
    if not ids:
        return []
    scores = _fuse(db, ids, semantic_ids, keyword_ids)
    chunks = db.scalars(select(KnowledgeChunk).where(KnowledgeChunk.id.in_(ids))).all()
    usable = []
    for chunk in chunks:
        if substantive_only and (chunk.block_type == BLOCK_TOC
                                 or (chunk.page_type == "body" and is_toc_entry(chunk.text)
                                     and len(chunk.text.strip()) < MIN_CHUNK_CHARS)):
            continue
        usable.append(chunk)
    usable.sort(key=lambda chunk: (-scores.get(chunk.id, 0.0), chunk.id))
    return [{**_row_dict(chunk), "rank": float(scores.get(chunk.id, 0.0))} for chunk in usable[:limit]]


def _page_diverse(chunks: list[dict], limit: int, *, page_limit: int, section_cap: int | None):
    """Spread the first pass over pages, then allow a bounded number per section."""
    picked: list[dict] = []
    per_page: dict[tuple, int] = {}
    per_section: dict[tuple, int] = {}
    remaining = []
    for chunk in chunks:
        page_key = (chunk.get("document_id"), chunk.get("page"))
        section = (chunk.get("document_id"), str(chunk.get("section") or "")) if chunk.get("section") else None
        if len(picked) < limit and per_page.get(page_key, 0) < page_limit:
            if section_cap and section and per_section.get(section, 0) >= section_cap:
                remaining.append(chunk)
                continue
            picked.append(chunk)
            per_page[page_key] = per_page.get(page_key, 0) + 1
            if section:
                per_section[section] = per_section.get(section, 0) + 1
        else:
            remaining.append(chunk)
    for chunk in remaining:
        if len(picked) >= limit:
            break
        picked.append(chunk)
    return picked[:limit]


def retrieve_targets(db, *, course_version_id: int, project_id: int, targets: list[dict], category=None,
                     categories=None, limit=SCORE_EVIDENCE_LIMIT, embedder=embed, include_public=False,
                     page_limit=DEFAULT_PAGE_LIMIT, section_cap=DEFAULT_SECTION_CAP,
                     per_target_limit: int | None = None, navigator=None, retrieve_fn=retrieve):
    """Search once per required-evidence target and merge with page-level diversity.

    Every target is searched separately with its own query and its own allowed page
    types, so a contents page permitted for a structural target can never flood a
    content target's evidence. A target counts as unmet when its *own* search found
    nothing admissible, which is what decides a teacher review.
    """
    merged: list[dict] = []
    seen: dict[int, dict] = {}
    per_target: dict[str, list[dict]] = {}
    for target in targets:
        allowed = target.get("allowed_page_types")
        raw_hits = retrieve_fn(db, course_version_id=course_version_id, project_id=project_id, category=category,
                               categories=categories, query=target.get("query") or target.get("label") or "",
                               limit=per_target_limit or max(limit, 8), embedder=embedder,
                               include_public=include_public, allowed_page_types=allowed,
                               substantive_only=target.get("use") != "structure",
                               semantic_pool=SEMANTIC_POOL, keyword_pool=KEYWORD_POOL)
        hits = [normalize_hit(hit) for hit in (raw_hits or [])]
        if target.get("use") == "content" and not hits and navigator is not None:
            # The contents list may only be used to navigate: its headings are matched
            # against headings that really occur in body text, and those sections and
            # their neighbouring pages are searched. The quotation still lands on body.
            hits = [dict(hit, via="toc_navigation") for hit in (navigator(target) or [])]
        per_target[target["target_id"]] = hits
        for hit in hits:
            if hit["chunk_id"] not in seen:
                item = {**hit, "target_id": target["target_id"], "target_ids": [target["target_id"]],
                        "target_label": target["label"], "target_use": target["use"]}
                seen[hit["chunk_id"]] = item
                merged.append(item)
            else:
                seen[hit["chunk_id"]]["target_ids"].append(target["target_id"])
    merged.sort(key=lambda chunk: -chunk["rank"])
    # Reserve coverage before filling the remaining budget by rank. One shared
    # passage may cover several targets, so it consumes only one slot.
    selected = []
    covered = set()
    for target in targets:
        if target["target_id"] in covered or len(selected) >= limit:
            continue
        candidates = [chunk for chunk in merged if target["target_id"] in chunk["target_ids"]]
        if candidates:
            selected.append(candidates[0])
            covered.update(candidates[0]["target_ids"])
    chosen = {chunk["chunk_id"] for chunk in selected}
    remaining = [chunk for chunk in merged if chunk["chunk_id"] not in chosen]
    selected.extend(_page_diverse(remaining, limit - len(selected), page_limit=page_limit, section_cap=section_cap))
    covered = {target_id for chunk in selected for target_id in chunk["target_ids"]}
    uncovered = [target for target in targets if target["target_id"] not in covered]
    return {"evidence": selected, "per_target": per_target, "uncovered_targets": uncovered}


def body_heading_index(db, *, course_version_id: int, project_id: int, category: str | None = None):
    """Sections present in real body text, used to navigate from a contents page."""
    filters = [KnowledgeChunk.course_version_id == course_version_id, KnowledgeChunk.project_id == project_id,
               KnowledgeChunk.page_type == "body", KnowledgeChunk.block_type != "legacy",
               KnowledgeChunk.document_id.in_(current_document_ids())]
    if category:
        filters.append(KnowledgeChunk.document_type == category)
    rows = db.execute(select(KnowledgeChunk.document_id, KnowledgeChunk.page, KnowledgeChunk.section,
                             KnowledgeChunk.chapter).where(*filters).distinct()).all()
    index: dict[str, list[dict]] = {}
    for document_id, page, section, chapter in rows:
        for value in (section, chapter):
            key = normalize_heading(value or "")
            if key:
                index.setdefault(key, []).append({"document_id": document_id, "page": page,
                                                  "section": section, "chapter": chapter})
    return index