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
512 lines
20 KiB
Python
512 lines
20 KiB
Python
"""Book-level orchestration for LLM judging and promotion of candidate phrases."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import re
|
|
from time import perf_counter
|
|
from typing import TYPE_CHECKING
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from python.ebook_search.llm_interface import request_chat_completion
|
|
from python.ebook_search.protected_phrases.extraction import (
|
|
candidate_source_names,
|
|
get_sample_contexts,
|
|
is_junk_phrase,
|
|
is_most_common_word_phrase,
|
|
score_candidate,
|
|
)
|
|
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
|
|
from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult
|
|
from python.ebook_search.protected_phrases.store import (
|
|
count_protected_phrases,
|
|
count_unjudged_candidates,
|
|
load_book_text,
|
|
load_candidates_for_judgment,
|
|
phrase_candidate_from_row,
|
|
save_candidate_to_db,
|
|
upsert_protected_phrase,
|
|
)
|
|
from python.ebook_search.protected_phrases.text_normalization import normalize_text
|
|
from python.orm.richie import EbookSource
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
|
from python.orm.richie import EbookProtectedPhrase
|
|
|
|
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def judge_candidate_phrases_for_books(
|
|
engine: AsyncEngine,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
source_ids: Sequence[int] | None = None,
|
|
) -> PhraseJudgmentBackfillResult:
|
|
"""Judge candidate phrases for books, fanning LLM calls out across books and phrases.
|
|
|
|
Up to ``phrase_judge_book_workers`` books are judged at once, and within each book candidates
|
|
are judged in concurrent chunks of ``phrase_judge_phrase_workers``. Each book uses its own
|
|
short-lived sessions for reads and writes; no database connection is held while LLM calls are
|
|
in flight. For a pseudo-single-threaded run (solo testing, debugging), set both worker
|
|
settings to 1.
|
|
|
|
Args:
|
|
engine (AsyncEngine): Engine used to open one session per book.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration.
|
|
source_ids (Sequence[int] | None): Books to judge; ``None`` judges every indexed book.
|
|
|
|
Returns:
|
|
PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates,
|
|
protected phrases, and mentions.
|
|
"""
|
|
if source_ids is None:
|
|
async with AsyncSession(engine) as session:
|
|
source_ids = list((await session.scalars(select(EbookSource.id).order_by(EbookSource.id))).all())
|
|
books_seen = len(source_ids)
|
|
book_workers = max(1, config.phrase_judge_book_workers)
|
|
phrase_workers = max(1, config.phrase_judge_phrase_workers)
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s "
|
|
"confidence_threshold=%.2f",
|
|
books_seen,
|
|
book_workers,
|
|
phrase_workers,
|
|
config.protected_phrase_confidence_threshold,
|
|
)
|
|
|
|
book_semaphore = asyncio.Semaphore(book_workers)
|
|
max_connections = book_workers * phrase_workers
|
|
limits = httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections)
|
|
async with httpx.AsyncClient(limits=limits) as client:
|
|
outcomes = await asyncio.gather(
|
|
*(judge_one_book_async(engine, source_id, config, client, book_semaphore) for source_id in source_ids)
|
|
)
|
|
|
|
result = PhraseJudgmentBackfillResult(
|
|
books_seen=books_seen,
|
|
books_judged=sum(1 for outcome in outcomes if outcome.committed),
|
|
books_failed=sum(1 for outcome in outcomes if outcome.failed),
|
|
candidates_judged=sum(outcome.judged for outcome in outcomes),
|
|
protected_phrases=sum(outcome.protected for outcome in outcomes),
|
|
phrase_mentions=sum(outcome.mentions for outcome in outcomes),
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s "
|
|
"candidates_judged=%s protected=%s mentions=%s",
|
|
result.books_seen,
|
|
result.books_judged,
|
|
result.books_failed,
|
|
result.candidates_judged,
|
|
result.protected_phrases,
|
|
result.phrase_mentions,
|
|
)
|
|
return result
|
|
|
|
|
|
async def judge_one_book_async(
|
|
engine: AsyncEngine,
|
|
source_id: int,
|
|
config: EbookSearchConfig,
|
|
client: httpx.AsyncClient,
|
|
book_semaphore: asyncio.Semaphore,
|
|
) -> BookJudgmentResult:
|
|
"""Judge one book concurrently and persist the outcome, honoring the book-level limit.
|
|
|
|
Args:
|
|
engine (AsyncEngine): Engine used to open the book's read and write sessions.
|
|
source_id (int): Book to judge candidates for.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
client (httpx.AsyncClient): Shared async client for LLM calls.
|
|
book_semaphore (asyncio.Semaphore): Caps how many books judge at once.
|
|
|
|
Returns:
|
|
BookJudgmentResult: The book's judgment outcome.
|
|
"""
|
|
async with book_semaphore:
|
|
try:
|
|
prepared = await prepare_book_judgment(engine, source_id, config)
|
|
if prepared is None:
|
|
return BookJudgmentResult()
|
|
work_items, target_remaining = prepared
|
|
judged = await judge_book_candidates_async(client, config, source_id, work_items, target_remaining)
|
|
if not judged:
|
|
return BookJudgmentResult()
|
|
return await persist_book_judgments(engine, source_id, config, judged)
|
|
except Exception:
|
|
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
|
|
return BookJudgmentResult(failed=True)
|
|
|
|
|
|
async def prepare_book_judgment(
|
|
engine: AsyncEngine,
|
|
source_id: int,
|
|
config: EbookSearchConfig,
|
|
) -> tuple[list[tuple[int, PhraseCandidate]], int | None] | None:
|
|
"""Load one book's candidates to judge, with sample contexts, on a short-lived read session.
|
|
|
|
Args:
|
|
engine (AsyncEngine): Engine used to open the read session.
|
|
source_id (int): Book to load candidates for.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
|
|
Returns:
|
|
tuple[list[tuple[int, PhraseCandidate]], int | None] | None: Candidate rows paired with
|
|
in-memory candidates and the remaining protected-phrase target, or ``None`` when the book
|
|
has nothing to judge.
|
|
"""
|
|
judgment_limit = config.protected_phrase_llm_candidates_per_book
|
|
if judgment_limit <= 0:
|
|
return None
|
|
async with AsyncSession(engine) as session:
|
|
if not await count_unjudged_candidates(session, source_id, config):
|
|
logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id)
|
|
return None
|
|
existing_protected = await count_protected_phrases(session, source_id)
|
|
target_remaining: int | None = None
|
|
if config.phrase_target_protected_per_book > 0:
|
|
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
|
|
if target_remaining == 0:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_skipped_target_met source_id=%s existing_protected=%s target=%s",
|
|
source_id,
|
|
existing_protected,
|
|
config.phrase_target_protected_per_book,
|
|
)
|
|
return None
|
|
book_text = await load_book_text(session, source_id)
|
|
if not book_text:
|
|
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
|
|
return None
|
|
normalized_book_text = normalize_text(book_text)
|
|
# Stored rows may predate the current junk filters and score weights, so re-filter and
|
|
# rescore every unjudged row here instead of trusting the persisted candidate_score.
|
|
rows = await load_candidates_for_judgment(session, source_id, config)
|
|
scored_items: list[tuple[int, PhraseCandidate]] = []
|
|
skipped_junk = 0
|
|
for row in rows:
|
|
candidate = phrase_candidate_from_row(row)
|
|
if is_junk_phrase(candidate.phrase_norm.split()):
|
|
skipped_junk += 1
|
|
continue
|
|
candidate.candidate_score = score_candidate(candidate, config)
|
|
scored_items.append((row.id, candidate))
|
|
scored_items.sort(key=lambda item: item[1].candidate_score, reverse=True)
|
|
work_items = scored_items[:judgment_limit]
|
|
for _, candidate in work_items:
|
|
candidate.sample_contexts = candidate.sample_contexts or get_sample_contexts(
|
|
normalized_book_text, candidate.phrase_norm
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s "
|
|
"unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s",
|
|
source_id,
|
|
len(work_items),
|
|
skipped_junk,
|
|
len(rows),
|
|
existing_protected,
|
|
target_remaining,
|
|
judgment_limit,
|
|
)
|
|
return work_items, target_remaining
|
|
|
|
|
|
async def judge_book_candidates_async(
|
|
client: httpx.AsyncClient,
|
|
config: EbookSearchConfig,
|
|
source_id: int,
|
|
work_items: list[tuple[int, PhraseCandidate]],
|
|
target_remaining: int | None,
|
|
) -> list[tuple[int, PhraseCandidate, LLMJudgment, bool]]:
|
|
"""Judge a book's candidates in concurrent chunks, stopping once the target is reached.
|
|
|
|
Promotion decisions are made in memory so judging can stop early without any database writes.
|
|
|
|
Args:
|
|
client (httpx.AsyncClient): Shared async client for LLM calls.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
source_id (int): Book being judged, for logging.
|
|
work_items (list[tuple[int, PhraseCandidate]]): Candidate row ids paired with candidates,
|
|
in best-first score order.
|
|
target_remaining (int | None): Remaining protected-phrase target, or ``None`` for no cap.
|
|
|
|
Returns:
|
|
list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: Judged rows with their judgment and
|
|
whether each should be promoted.
|
|
"""
|
|
chunk_size = max(1, config.phrase_judge_phrase_workers)
|
|
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]] = []
|
|
promoted = 0
|
|
for start in range(0, len(work_items), chunk_size):
|
|
chunk = work_items[start : start + chunk_size]
|
|
judgments = await asyncio.gather(*(judge_candidate_async(client, config, candidate) for _, candidate in chunk))
|
|
for (candidate_id, candidate), judgment in zip(chunk, judgments, strict=True):
|
|
promote = (target_remaining is None or promoted < target_remaining) and should_protect_judged_candidate(
|
|
candidate, judgment, source_id, config, candidate_id=candidate_id
|
|
)
|
|
if promote:
|
|
promoted += 1
|
|
judged.append((candidate_id, candidate, judgment, promote))
|
|
if target_remaining is not None and promoted >= target_remaining:
|
|
break
|
|
return judged
|
|
|
|
|
|
async def judge_candidate_async(
|
|
client: httpx.AsyncClient,
|
|
config: EbookSearchConfig,
|
|
candidate: PhraseCandidate,
|
|
) -> LLMJudgment:
|
|
"""Judge one candidate with the LLM over the shared async client.
|
|
|
|
Args:
|
|
client (httpx.AsyncClient): Shared async client for LLM calls.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
candidate (PhraseCandidate): Candidate to judge.
|
|
|
|
Returns:
|
|
LLMJudgment: The parsed judgment.
|
|
"""
|
|
content = await request_chat_completion(client, config, build_judge_messages(candidate))
|
|
return parse_llm_judgment(content, config)
|
|
|
|
|
|
async def persist_book_judgments(
|
|
engine: AsyncEngine,
|
|
source_id: int,
|
|
config: EbookSearchConfig,
|
|
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]],
|
|
) -> BookJudgmentResult:
|
|
"""Persist one book's judgments and promotions in a single committed transaction.
|
|
|
|
Args:
|
|
engine (AsyncEngine): Engine used to open the write session.
|
|
source_id (int): Book being persisted.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
judged (list[tuple[int, PhraseCandidate, LLMJudgment, bool]]): Judged candidates with their
|
|
judgment and promotion flag.
|
|
|
|
Returns:
|
|
BookJudgmentResult: The book's committed counts, or a failed result on error.
|
|
"""
|
|
book_started_at = perf_counter()
|
|
async with AsyncSession(engine, expire_on_commit=False) as session:
|
|
try:
|
|
protected: list[EbookProtectedPhrase] = []
|
|
for candidate_id, candidate, judgment, promote in judged:
|
|
candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment)
|
|
if promote:
|
|
protected.append(
|
|
await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row)
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r "
|
|
"keep=%s confidence=%.3f category=%r promoted=%s",
|
|
source_id,
|
|
candidate_id,
|
|
candidate.phrase_norm,
|
|
judgment.keep,
|
|
judgment.confidence,
|
|
judgment.category,
|
|
promote,
|
|
)
|
|
await session.flush()
|
|
mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
logger.exception("ebook_candidate_phrase_judgment_book_persist_failed source_id=%s", source_id)
|
|
return BookJudgmentResult(failed=True)
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
|
|
"duration_ms=%.1f",
|
|
source_id,
|
|
len(judged),
|
|
len(protected),
|
|
mentions,
|
|
(perf_counter() - book_started_at) * 1000,
|
|
)
|
|
return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True)
|
|
|
|
|
|
def should_protect_judged_candidate(
|
|
candidate: PhraseCandidate,
|
|
judgment: LLMJudgment,
|
|
book_id: int,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
candidate_id: int,
|
|
) -> bool:
|
|
"""Report whether a judged candidate qualifies to become a protected phrase.
|
|
|
|
Args:
|
|
candidate (PhraseCandidate): In-memory candidate that was judged.
|
|
judgment (LLMJudgment): Judge decision for the candidate.
|
|
book_id (int): Book the candidate belongs to, for logging.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
candidate_id (int): Stored candidate row id the judgment came from, for logging.
|
|
|
|
Returns:
|
|
bool: True when the judged candidate should be promoted to a protected phrase.
|
|
"""
|
|
if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold:
|
|
return False
|
|
accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text)
|
|
accepted_tokens = accepted_norm.split()
|
|
accepted_token_count = len(accepted_tokens)
|
|
if accepted_token_count < config.phrase_min_tokens:
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s "
|
|
"phrase=%r canonical=%r token_count=%s min_tokens=%s",
|
|
book_id,
|
|
candidate_id,
|
|
candidate.phrase_norm,
|
|
accepted_norm,
|
|
accepted_token_count,
|
|
config.phrase_min_tokens,
|
|
)
|
|
return False
|
|
if is_most_common_word_phrase(accepted_tokens):
|
|
logger.info(
|
|
"ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s "
|
|
"phrase=%r canonical=%r",
|
|
book_id,
|
|
candidate_id,
|
|
candidate.phrase_norm,
|
|
accepted_norm,
|
|
)
|
|
return False
|
|
return True
|
|
|
|
|
|
def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]:
|
|
"""Build the chat messages used to judge one candidate phrase.
|
|
|
|
Args:
|
|
candidate (PhraseCandidate): Candidate to describe for the judge.
|
|
|
|
Returns:
|
|
list[dict[str, str]]: OpenAI-style system and user messages.
|
|
"""
|
|
payload = {
|
|
"phrase": candidate.phrase_norm,
|
|
"token_count": candidate.token_count,
|
|
"sources": candidate_source_names(candidate),
|
|
"raw_count": candidate.raw_count,
|
|
"chapter_count": candidate.chapter_count,
|
|
"contexts": candidate.sample_contexts,
|
|
}
|
|
return [
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"Judge whether a candidate phrase from a book should be protected for RAG retrieval. "
|
|
"Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, "
|
|
"unstable fragments, and phrases kept only because they are frequent. Keep people, places, "
|
|
"organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, "
|
|
"named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, "
|
|
"canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason."
|
|
),
|
|
},
|
|
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
|
|
]
|
|
|
|
|
|
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
|
|
"""Parse and validate an LLM phrase-judge response.
|
|
|
|
Args:
|
|
content (str): Raw model response text.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings supplying nesting defaults.
|
|
|
|
Returns:
|
|
LLMJudgment: The parsed and validated judgment.
|
|
|
|
Raises:
|
|
TypeError: If the decoded JSON body is not an object.
|
|
"""
|
|
body = json.loads(extract_json_object(content))
|
|
if not isinstance(body, dict):
|
|
msg = "LLM phrase judge response is not a JSON object"
|
|
raise TypeError(msg)
|
|
|
|
aliases = body.get("aliases", ())
|
|
if not isinstance(aliases, list | tuple):
|
|
aliases = ()
|
|
return LLMJudgment(
|
|
keep=bool(body.get("keep", False)),
|
|
canonical=optional_text(body.get("canonical")),
|
|
category=optional_text(body.get("category")),
|
|
aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()),
|
|
confidence=clamped_float(body.get("confidence"), default=0.0),
|
|
importance=clamped_float(body.get("importance"), default=0.5),
|
|
allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)),
|
|
suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)),
|
|
reason=optional_text(body.get("reason")),
|
|
)
|
|
|
|
|
|
def extract_json_object(content: str) -> str:
|
|
"""Extract a JSON object from plain or fenced model output.
|
|
|
|
Args:
|
|
content (str): Raw model response text.
|
|
|
|
Returns:
|
|
str: The substring spanning the first JSON object.
|
|
|
|
Raises:
|
|
ValueError: If no JSON object is found in the response.
|
|
"""
|
|
stripped = content.strip()
|
|
if stripped.startswith("{") and stripped.endswith("}"):
|
|
return stripped
|
|
match = JSON_OBJECT_RE.search(stripped)
|
|
if match is None:
|
|
msg = "LLM phrase judge response did not contain a JSON object"
|
|
raise ValueError(msg)
|
|
return match.group(0)
|
|
|
|
|
|
def optional_text(value: object) -> str | None:
|
|
"""Return stripped text for a nullable JSON value.
|
|
|
|
Args:
|
|
value (object): Decoded JSON value that may or may not be a string.
|
|
|
|
Returns:
|
|
str | None: The stripped string, or ``None`` when it is not a non-empty string.
|
|
"""
|
|
if not isinstance(value, str):
|
|
return None
|
|
stripped = value.strip()
|
|
return stripped or None
|
|
|
|
|
|
def clamped_float(value: object, *, default: float) -> float:
|
|
"""Coerce a JSON number into the 0.0 to 1.0 range.
|
|
|
|
Args:
|
|
value (object): Decoded JSON value that may or may not be a number.
|
|
default (float): Fallback returned when ``value`` is not numeric.
|
|
|
|
Returns:
|
|
float: The value clamped to ``[0.0, 1.0]``, or ``default`` when non-numeric.
|
|
"""
|
|
if not isinstance(value, int | float):
|
|
return default
|
|
return min(max(float(value), 0.0), 1.0)
|