"""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), ), ), )