feat: add chunk BM25 and reproducible RRF retrieval

2026-10-16 13:57:35 +08:00
parent ccf1900996
commit bb273ec724
13 changed files with 797 additions and 30 deletions
@@ -7,13 +7,23 @@ from sqlalchemy.orm import Session
 from app.db.session import get_db
 from app.models import RetrievalRun
 from app.schemas.catalog import Page
+from app.schemas.hybrid import (
+    ChunkKeywordRequest,
+    ChunkKeywordResponse,
+    HybridRequest,
+    HybridResponse,
+)
 from app.schemas.retrieval import RetrievalRequest, RetrievalResponse
 from app.schemas.vectors import VectorRetrievalResponse
 from app.services.auth import require_admin
+from app.services.embedding_runtime import get_embedding_provider
+from app.services.hybrid import retrieve_chunks, retrieve_hybrid
 from app.services.retrieval import retrieve

 router = APIRouter(prefix="/retrieval", tags=["retrieval"])

+RunResponse = RetrievalResponse | VectorRetrievalResponse | ChunkKeywordResponse | HybridResponse
+

@router.post("", response_model=RetrievalResponse)
 def keyword_retrieval(payload: RetrievalRequest, db: Session = Depends(get_db)):
@@ -21,9 +31,23 @@ def keyword_retrieval(payload: RetrievalRequest, db: Session = Depends(get_db)):
    return retrieve(db, payload)


+@router.post("/chunks", response_model=ChunkKeywordResponse)
+def chunk_keyword_retrieval(payload: ChunkKeywordRequest, db: Session = Depends(get_db)):
+    return retrieve_chunks(db, payload)
+
+
+@router.post("/hybrid", response_model=HybridResponse)
+def hybrid_retrieval(
+    payload: HybridRequest,
+    db: Session = Depends(get_db),
+    provider=Depends(get_embedding_provider),
+):
+    return retrieve_hybrid(db, payload, provider)
+
+
@router.get(
    "/runs",
-    response_model=Page[RetrievalResponse | VectorRetrievalResponse],
+    response_model=Page[RunResponse],
    dependencies=[Depends(require_admin)],
 )
 def list_runs(
@@ -47,7 +71,7 @@ def list_runs(

@router.get(
    "/runs/{run_id}",
-    response_model=RetrievalResponse | VectorRetrievalResponse,
+    response_model=RunResponse,
    dependencies=[Depends(require_admin)],
 )
 def read_run(run_id: UUID, db: Session = Depends(get_db)):
@@ -0,0 +1,95 @@
+"""Chunk-level lexical and hybrid contracts; the item-level baseline stays distinct."""
+
+from datetime import datetime
+from typing import Literal
+from uuid import UUID
+
+from pydantic import BaseModel, Field, model_validator
+
+from app.schemas.catalog import Category, Input
+from app.schemas.indexing import IndexConfig
+from app.schemas.retrieval import RetrievalConfig
+from app.schemas.vectors import SpaceRead, VectorHit
+
+
+class ChunkKeywordRequest(Input):
+    query: str = Field(min_length=1, max_length=500)
+    mode: Literal["keyword_chunks"] = "keyword_chunks"
+    top_k: int = Field(5, ge=1, le=20)
+    category: Category | None = None
+
+
+class HybridRequest(ChunkKeywordRequest):
+    mode: Literal["hybrid"] = "hybrid"
+    candidate_k: int = Field(50, ge=1, le=200)
+
+    @model_validator(mode="after")
+    def sufficient_candidates(self):
+        if self.candidate_k < self.top_k:
+            raise ValueError("candidate_k 不得小于 top_k")
+        return self
+
+
+class ChunkRetrievalConfig(BaseModel):
+    ranking_unit: Literal["chunk"] = "chunk"
+    lexical: RetrievalConfig
+    chunking: IndexConfig
+    tie_break: Literal["knowledge_id,chunk_ordinal,chunk_id"] = (
+        "knowledge_id,chunk_ordinal,chunk_id"
+    )
+
+
+class HybridConfig(ChunkRetrievalConfig):
+    scorer: Literal["rrf-v1"] = "rrf-v1"
+    rrf_k: Literal[60] = 60
+    weights: tuple[Literal[1], Literal[1]] = (1, 1)
+    candidate_k: int
+    vector_scorer: Literal["pgvector-cosine-exact-v1"] = "pgvector-cosine-exact-v1"
+    space: SpaceRead
+
+
+class RankedCandidate(BaseModel):
+    chunk_id: UUID
+    rank: int
+    score: float
+
+
+class HybridHit(VectorHit):
+    lexical_rank: int | None
+    lexical_score: float | None
+    vector_rank: int | None
+    vector_score: float | None
+    lexical_contribution: float
+    vector_contribution: float
+
+
+class ChunkKeywordResponse(BaseModel):
+    run_id: UUID
+    query: str
+    mode: Literal["keyword_chunks"] = "keyword_chunks"
+    status: Literal["ok", "no_evidence"]
+    top_k: int
+    category: str | None
+    query_terms: list[str]
+    text_index_version_id: UUID
+    text_generation: int
+    catalog_revision: int
+    corpus_fingerprint: str
+    corpus_size: int
+    total_hits: int
+    results: list[VectorHit]
+    elapsed_ms: float
+    created_at: datetime
+    config: ChunkRetrievalConfig
+
+
+class HybridResponse(ChunkKeywordResponse):
+    mode: Literal["hybrid"] = "hybrid"
+    vector_version_id: UUID
+    generation: int
+    lexical_total_hits: int
+    vector_total_hits: int
+    lexical_candidates: list[RankedCandidate]
+    vector_candidates: list[RankedCandidate]
+    results: list[HybridHit]
+    config: HybridConfig
@@ -0,0 +1,205 @@
+"""Bounded chunk BM25 and equal-weight reciprocal rank fusion on one publication."""
+
+from collections import Counter
+from time import perf_counter
+from uuid import uuid4
+
+from fastapi import HTTPException
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from app.models.catalog import utc_now
+from app.models.indexing import Chunk
+from app.models.retrieval import RetrievalRun
+from app.schemas.hybrid import (
+    ChunkKeywordRequest,
+    ChunkKeywordResponse,
+    ChunkRetrievalConfig,
+    HybridConfig,
+    HybridHit,
+    HybridRequest,
+    HybridResponse,
+    RankedCandidate,
+)
+from app.schemas.vectors import SpaceRead, VectorHit
+from app.services.indexing import locked_state
+from app.services.retrieval import CONFIG, bm25, tokenize
+from app.services.vectors import (
+    checked_chunks,
+    current_text,
+    encode,
+    rank_vectors,
+    retrieval_scope,
+)
+
+RRF_K = 60
+CHUNK_LEXICAL_CONFIG = CONFIG.model_copy(update={"scorer": "bm25-chunk-body-v1", "title_weight": 0})
+
+
+def chunk_key(chunk):
+    return chunk.knowledge_id, chunk.ordinal, chunk.id
+
+
+def chunks_in_scope(db, text_version, category):
+    statement = select(Chunk).where(Chunk.index_version_id == text_version.id)
+    if category is not None:
+        statement = statement.where(Chunk.provenance["category"].as_string() == category)
+    return list(db.scalars(statement.order_by(Chunk.knowledge_id, Chunk.ordinal, Chunk.id)))
+
+
+def lexical_ranking(chunks, terms):
+    counters = {chunk.id: Counter(tokenize(chunk.text)) for chunk in chunks}
+    scores = bm25(list(counters.values()), terms)
+    ranked = sorted(
+        ((chunk, score) for chunk, score in zip(chunks, scores) if score > 0),
+        key=lambda pair: (-pair[1], chunk_key(pair[0])),
+    )
+    return ranked, counters
+
+
+def candidate_list(rows):
+    return [
+        RankedCandidate(chunk_id=chunk.id, rank=rank, score=score)
+        for rank, (chunk, score) in enumerate(rows, 1)
+    ]
+
+
+def fuse(lexical, vector):
+    """Ranks start at one; an absent candidate contributes zero, never a padded rank."""
+    scores = {}
+    for leg in (lexical, vector):
+        for candidate in leg:
+            key = str(candidate.chunk_id)
+            scores[key] = scores.get(key, 0.0) + 1.0 / (RRF_K + candidate.rank)
+    return scores
+
+
+def hit_fields(chunk, terms, counters):
+    return dict(
+        **chunk.provenance,
+        chunk_id=chunk.id,
+        chunk_ordinal=chunk.ordinal,
+        snippet=chunk.text,
+        snippet_start=chunk.start_offset,
+        snippet_end=chunk.end_offset,
+        matched_terms=[term for term in terms if counters[chunk.id][term]],
+    )
+
+
+def save_run(db, response):
+    db.add(
+        RetrievalRun(
+            id=str(response.run_id),
+            query=response.query,
+            mode=response.mode,
+            corpus_fingerprint=response.corpus_fingerprint,
+            response=response.model_dump(mode="json"),
+            created_at=response.created_at,
+        )
+    )
+    db.commit()
+    return response
+
+
+def response_fields(payload, state, text_version, chunks, terms, total, started):
+    return dict(
+        run_id=uuid4(),
+        query=payload.query,
+        status="ok" if total else "no_evidence",
+        top_k=payload.top_k,
+        category=payload.category,
+        query_terms=terms,
+        text_index_version_id=text_version.id,
+        text_generation=state.generation,
+        catalog_revision=state.catalog_revision,
+        corpus_fingerprint=text_version.corpus_fingerprint,
+        corpus_size=len(chunks),
+        total_hits=total,
+        elapsed_ms=round((perf_counter() - started) * 1000, 3),
+        created_at=utc_now(),
+    )
+
+
+def retrieve_chunks(db: Session, payload: ChunkKeywordRequest):
+    started = perf_counter()
+    terms = list(dict.fromkeys(tokenize(payload.query)))
+    if not terms:
+        raise HTTPException(422, "问题中没有可检索词，请补充具体关键词")
+    state = locked_state(db)
+    text_version, corpus = current_text(db, state)
+    checked_chunks(db, text_version, corpus)
+    chunks = chunks_in_scope(db, text_version, payload.category)
+    ranked, counters = lexical_ranking(chunks, terms)
+    hits = [
+        VectorHit(**hit_fields(chunk, terms, counters), rank=rank, score=score)
+        for rank, (chunk, score) in enumerate(ranked[: payload.top_k], 1)
+    ]
+    return save_run(
+        db,
+        ChunkKeywordResponse(
+            **response_fields(payload, state, text_version, chunks, terms, len(ranked), started),
+            results=hits,
+            config=ChunkRetrievalConfig(lexical=CHUNK_LEXICAL_CONFIG, chunking=text_version.config),
+        ),
+    )
+
+
+def retrieve_hybrid(db: Session, payload: HybridRequest, provider):
+    started = perf_counter()
+    terms = list(dict.fromkeys(tokenize(payload.query)))
+    *_, token = retrieval_scope(db, provider)
+    db.commit()
+    query_vector = encode(provider, [payload.query], query=True)[0]
+    state, vector_state, text_version, version, space, current = retrieval_scope(db, provider)
+    if token != current:
+        raise HTTPException(409, "查询编码期间发布版本或审核依据已变化，请重试")
+    # Hold the shared catalog/publication locks until both rankings and the run are saved.
+    chunks = chunks_in_scope(db, text_version, payload.category)
+    ranked, counters = lexical_ranking(chunks, terms)
+    lexical = candidate_list(ranked[: payload.candidate_k])
+    vector_total, vector_rows = rank_vectors(
+        db, version, space, payload.category, payload.candidate_k, query_vector
+    )
+    vector = candidate_list(
+        [(chunk, max(-1.0, min(1.0, 1.0 - float(distance)))) for chunk, distance in vector_rows]
+    )
+    scores = fuse(lexical, vector)
+    by_id = {chunk.id: chunk for chunk in chunks}
+    ordered = sorted(scores, key=lambda key: (-scores[key], chunk_key(by_id[key])))
+    lexical_by_id = {str(c.chunk_id): c for c in lexical}
+    vector_by_id = {str(c.chunk_id): c for c in vector}
+    hits = []
+    for rank, key in enumerate(ordered[: payload.top_k], 1):
+        left, right = lexical_by_id.get(key), vector_by_id.get(key)
+        hits.append(
+            HybridHit(
+                **hit_fields(by_id[key], terms, counters),
+                rank=rank,
+                score=scores[key],
+                lexical_rank=left.rank if left else None,
+                lexical_score=left.score if left else None,
+                vector_rank=right.rank if right else None,
+                vector_score=right.score if right else None,
+                lexical_contribution=1.0 / (RRF_K + left.rank) if left else 0.0,
+                vector_contribution=1.0 / (RRF_K + right.rank) if right else 0.0,
+            )
+        )
+    return save_run(
+        db,
+        HybridResponse(
+            **response_fields(payload, state, text_version, chunks, terms, len(scores), started),
+            vector_version_id=version.id,
+            generation=vector_state.generation,
+            lexical_total_hits=len(ranked),
+            vector_total_hits=vector_total,
+            lexical_candidates=lexical,
+            vector_candidates=vector,
+            results=hits,
+            config=HybridConfig(
+                lexical=CHUNK_LEXICAL_CONFIG,
+                chunking=text_version.config,
+                candidate_k=payload.candidate_k,
+                space=SpaceRead.model_validate(space),
+            ),
+        ),
+    )
@@ -287,26 +287,7 @@ def retrieve_vectors(db: Session, payload: VectorRequest, provider: EmbeddingPro
    state, vector, text_version, version, space, current = retrieval_scope(db, provider)
    if token != current:
        raise HTTPException(409, "查询编码期间发布版本或审核依据已变化，请重试")
-    # Extension-qualified operator works even with an isolated business schema search_path.
-    distance = ChunkVector.embedding.op("OPERATOR(public.<=>)", return_type=Float)(query_vector)
-    statement = (
-        select(Chunk, distance.label("distance"))
-        .join(ChunkVector, ChunkVector.chunk_id == Chunk.id)
-        .where(
-            ChunkVector.version_id == version.id,
-            ChunkVector.space_id == space.id,
-            ChunkVector.text_index_version_id == text_version.id,
-            ChunkVector.dimension == space.dimension,
-        )
-    )
-    if payload.category is not None:
-        statement = statement.where(Chunk.provenance["category"].as_string() == payload.category)
-    total = db.scalar(select(func.count()).select_from(statement.subquery()))
-    rows = db.execute(
-        statement.order_by(distance, Chunk.knowledge_id, Chunk.ordinal, Chunk.id).limit(
-            payload.top_k
-        )
-    )
+    total, rows = rank_vectors(db, version, space, payload.category, payload.top_k, query_vector)
    hits = [
        VectorHit(
            **chunk.provenance,
@@ -351,3 +332,25 @@ def retrieve_vectors(db: Session, payload: VectorRequest, provider: EmbeddingPro
    )
    db.commit()
    return result
+
+
+def rank_vectors(db, version, space, category, limit, query_vector):
+    # Extension-qualified operator works even with an isolated business schema search_path.
+    distance = ChunkVector.embedding.op("OPERATOR(public.<=>)", return_type=Float)(query_vector)
+    statement = (
+        select(Chunk, distance.label("distance"))
+        .join(ChunkVector, ChunkVector.chunk_id == Chunk.id)
+        .where(
+            ChunkVector.version_id == version.id,
+            ChunkVector.space_id == space.id,
+            ChunkVector.text_index_version_id == version.text_index_version_id,
+            ChunkVector.dimension == space.dimension,
+        )
+    )
+    if category is not None:
+        statement = statement.where(Chunk.provenance["category"].as_string() == category)
+    total = db.scalar(select(func.count()).select_from(statement.subquery()))
+    rows = db.execute(
+        statement.order_by(distance, Chunk.knowledge_id, Chunk.ordinal, Chunk.id).limit(limit)
+    )
+    return total, list(rows)