"""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, delete, func, literal, or_, select from .config import get_settings from .models import KnowledgeChunk, KnowledgeDocument, KnowledgePage from .pages import (BLOCK_BODY, BLOCK_TOC, classify_document, is_toc_entry, normalize_heading, searchable_pages) INDEX_VERSION = 3 #: 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 _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``/``generate_series`` 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 [] total = len(ascii_tokens) + len(bigrams) score = None for token in ascii_tokens: hit = case((func.strpos(func.lower(KnowledgeChunk.text), literal(token)) > 0, 1.0), else_=0.0) score = hit if score is None else score + hit if bigrams: series = select(func.generate_series(1, func.greatest(func.length(KnowledgeChunk.text) - 1, 0)).label("i")).subquery("series") matches = (select(func.count()).select_from(series) .where(func.substr(KnowledgeChunk.text, series.c.i, 2).in_(bigrams)).scalar_subquery()) score = matches if score is None else score + matches 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 split_pages(pages: list[dict], *, size=650, overlap=100): """Split pages into searchable chunks carrying page type and section path.""" 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: 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): yield {"page": page["page_number"], "page_type": page_type, "block_type": BLOCK_BODY, "chapter": chapter, "section": section, "text": text_block[start:start + size]} 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): existing = db.scalar(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)) if existing and existing.source_hash == source_hash and existing.index_version == INDEX_VERSION and existing.status == "ready": 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("嵌入向量数量与知识片段数量不一致") if existing: db.execute(delete(KnowledgePage).where(KnowledgePage.document_id == existing.id)) db.execute(delete(KnowledgeChunk).where(KnowledgeChunk.document_id == existing.id)) doc = existing doc.document_type, doc.source_hash = category, source_hash else: 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) 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)) 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 _base_filters(*, course_version_id: int, project_id: int | None, category, categories, include_public: bool, public_only: bool, allowed_page_types, exclude_document_ids): filters = [KnowledgeChunk.course_version_id == course_version_id, _scope_filter(project_id, include_public, public_only)] 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 [] total = len(ascii_tokens) + len(bigrams) score = None for token in ascii_tokens: term = literal(token) hit = case((func.strpos(func.lower(KnowledgeChunk.text), term) > 0, 1.0), else_=0.0) score = hit if score is None else score + hit if bigrams: start = func.generate_series(1, func.greatest(func.length(KnowledgeChunk.text) - 1, 0)) series = select(start.label("i")).subquery("series") matches = (select(func.count()).select_from(series) .where(func.substr(KnowledgeChunk.text, series.c.i, 2).in_(bigrams)) .scalar_subquery()) score = matches if score is None else score + matches 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() hits = sum(1 for token in ascii_tokens if token in lower) 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): """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) 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[str, int] = {} remaining = [] for chunk in chunks: page_key = (chunk.get("document_id"), chunk.get("page")) section = str(chunk.get("section") or "") 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: set[int] = set() 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=embed, 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: seen.add(hit["chunk_id"]) merged.append({**hit, "target_id": target["target_id"], "target_label": target["label"], "target_use": target["use"]}) merged.sort(key=lambda chunk: -chunk["rank"]) selected = _page_diverse(merged, limit, page_limit=page_limit, section_cap=section_cap) uncovered = [target for target in targets if not per_target.get(target["target_id"])] 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"] 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 def page_text(source: dict) -> str: blocks = [str(source.get("text") or "").strip()] blocks.extend("表格: " + str(table.get("rows", table)) for table in source.get("tables", [])) return "\n".join(block for block in blocks if block) def relative_image_path(value: str | None) -> str | None: if not value: return None root = get_settings().processed_path.resolve() path = Path(value).resolve() if not path.is_relative_to(root): raise ValueError("页面图像路径不在处理目录内") return str(path.relative_to(root))