Richie
2706c4417d
Convert the ebook-search web app to async end to end and add concurrency to the protected-phrase extraction and judging pipeline so large books no longer block the event loop or the UI. ORM / infra: - Add get_async_postgres_engine and factor shared URL/connect_args building into build_postgres_url (reused by the sync and async engine builders) - Add async FastAPI session helpers (get_async_db, AsyncDbSession) with expire_on_commit=False to avoid implicit IO under asyncio App: - Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings, answer, rerank and LLM calls; convert handlers to async - Share a single httpx.AsyncClient in app state for LLM requests; size the connection pool for concurrent phrase-judging workers - Add judge_tasks: run per-book judging as tracked background tasks so a book already being judged isn't double-queued Protected phrases: - Add a process pool (pool.py) and worker-count config (extraction/judge book/phrase workers) to parallelize candidate generation and judging - Split admin actions into all/missing variants for generation and judging Config: - Add protected_phrase_extraction_workers, phrase_judge_book_workers, phrase_judge_phrase_workers
141 lines
4.5 KiB
Python
141 lines
4.5 KiB
Python
"""vLLM-backed optional reranking."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass, replace
|
|
from typing import TYPE_CHECKING
|
|
|
|
from python.ebook_search.llm_interface import request_rerank
|
|
|
|
if TYPE_CHECKING:
|
|
import httpx
|
|
|
|
from python.ebook_search.config import RerankConfig
|
|
from python.ebook_search.search import SearchResult
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RerankResult:
|
|
"""A relevance score for one candidate chunk."""
|
|
|
|
chunk_id: int
|
|
score: float
|
|
|
|
|
|
async def rerank_chunks(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
candidates: list[SearchResult],
|
|
config: RerankConfig,
|
|
) -> list[SearchResult]:
|
|
"""Rerank candidates with a vLLM rerank endpoint."""
|
|
if not candidates:
|
|
return []
|
|
|
|
logger.info(
|
|
"ebook_rerank_request_start base_url=%s model=%s candidates=%s",
|
|
config.base_url,
|
|
config.model,
|
|
len(candidates),
|
|
)
|
|
scores = await score_candidates(client, query, candidates, config)
|
|
results = sorted(
|
|
(
|
|
replace(
|
|
result,
|
|
score=final_rerank_score(result, scores[result.chunk_id].score, candidates, config),
|
|
rerank_score=scores[result.chunk_id].score,
|
|
)
|
|
for result in candidates
|
|
),
|
|
key=lambda result: result.score,
|
|
reverse=True,
|
|
)
|
|
logger.info(
|
|
"ebook_rerank_request_complete base_url=%s model=%s candidates=%s",
|
|
config.base_url,
|
|
config.model,
|
|
len(results),
|
|
)
|
|
return results
|
|
|
|
|
|
async def score_candidates(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
candidates: list[SearchResult],
|
|
config: RerankConfig,
|
|
) -> dict[int, RerankResult]:
|
|
"""Score candidate chunks with the configured rerank API."""
|
|
body = await request_rerank(client, query, [candidate.text for candidate in candidates], config)
|
|
if body is None:
|
|
return zero_rerank_scores(candidates)
|
|
|
|
scores = parse_vllm_scores(body, candidates)
|
|
for result in scores.values():
|
|
logger.debug("ebook_rerank_candidate_scored chunk_id=%s score=%s", result.chunk_id, result.score)
|
|
return scores
|
|
|
|
|
|
def parse_vllm_scores(body: object, candidates: list[SearchResult]) -> dict[int, RerankResult]:
|
|
"""Parse vLLM rerank scores into chunk-id keyed results."""
|
|
if not isinstance(body, dict):
|
|
logger.debug("ebook_rerank_response_not_object", extra={"response": body})
|
|
return zero_rerank_scores(candidates)
|
|
|
|
results = body.get("results") or body.get("data")
|
|
if not isinstance(results, list):
|
|
logger.debug("ebook_rerank_response_missing_results", extra={"response": body})
|
|
return zero_rerank_scores(candidates)
|
|
|
|
scores = zero_rerank_scores(candidates)
|
|
for item in results:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
index = item.get("index")
|
|
score = item.get("relevance_score", item.get("score"))
|
|
if not isinstance(index, int) or index < 0 or index >= len(candidates):
|
|
continue
|
|
if not isinstance(score, int | float):
|
|
continue
|
|
chunk_id = candidates[index].chunk_id
|
|
scores[chunk_id] = RerankResult(chunk_id=chunk_id, score=clamp_score(float(score)))
|
|
return scores
|
|
|
|
|
|
def zero_rerank_scores(candidates: list[SearchResult]) -> dict[int, RerankResult]:
|
|
"""Return zero relevance scores for all candidate chunks."""
|
|
return {candidate.chunk_id: RerankResult(chunk_id=candidate.chunk_id, score=0.0) for candidate in candidates}
|
|
|
|
|
|
def clamp_score(score: float) -> float:
|
|
"""Clamp a rerank score into the supported 0.0 to 1.0 range."""
|
|
return min(max(score, 0.0), 1.0)
|
|
|
|
|
|
def final_rerank_score(
|
|
result: SearchResult,
|
|
rerank_score: float,
|
|
candidates: list[SearchResult],
|
|
config: RerankConfig,
|
|
) -> float:
|
|
"""Combine rerank relevance with normalized hybrid retrieval evidence."""
|
|
return (config.score_weight * rerank_score) + (config.hybrid_weight * normalized_hybrid_score(result, candidates))
|
|
|
|
|
|
def normalized_hybrid_score(result: SearchResult, candidates: list[SearchResult]) -> float:
|
|
"""Normalize a candidate hybrid score against the rerank candidate set."""
|
|
hybrid_scores = [
|
|
candidate.fused_score if candidate.fused_score is not None else candidate.score for candidate in candidates
|
|
]
|
|
low = min(hybrid_scores)
|
|
high = max(hybrid_scores)
|
|
if high == low:
|
|
return 1.0
|
|
|
|
score = result.fused_score if result.fused_score is not None else result.score
|
|
return (score - low) / (high - low)
|