Files
dotfiles/python/ebook_search/api/judge_tasks.py
T

128 lines
4.8 KiB
Python

"""Background phrase-judging tasks for the web app.
Judging a book sends one LLM request per candidate phrase, which can take minutes, so it must
not run inside the request where it would block the UI. Judgments run as async FastAPI
background tasks, awaited on the event loop after the response is sent, and are tracked per
book in app state so a second judge request for a book that is already being judged is
rejected instead of doubling the work.
State is loop-confined: every read and mutation happens on the event loop (async route
handlers and async background tasks) and no critical section contains an ``await``, so each
mutation is atomic per loop iteration and no locking is needed.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books
if TYPE_CHECKING:
from fastapi import BackgroundTasks, FastAPI
logger = logging.getLogger(__name__)
@dataclass
class JudgeTaskState:
"""Running book judgments and last outcome messages, keyed by book id."""
running_book_ids: set[int] = field(default_factory=set)
outcome_messages: dict[int, str] = field(default_factory=dict)
def get_judge_task_state(app: FastAPI) -> JudgeTaskState:
"""Return the app's judge task state, creating it on first use.
Args:
app (FastAPI): App whose state holds the judge task registry.
Returns:
JudgeTaskState: The shared judge task state for this app.
"""
state = getattr(app.state, "judge_tasks", None)
if state is None:
state = JudgeTaskState()
app.state.judge_tasks = state
return state
def start_book_phrase_judgment(app: FastAPI, background_tasks: BackgroundTasks, source_id: int) -> bool:
"""Queue judging of one book's candidate phrases as a FastAPI background task.
The book is claimed before the response returns, so a repeated judge request cannot queue
a second run while one is pending or running.
Args:
app (FastAPI): App supplying the engine, config, and judge task state.
background_tasks (BackgroundTasks): Request's background tasks to queue the judgment on.
source_id (int): Book to judge candidates for.
Returns:
bool: True when a judgment was queued, False when one is already running for this book.
"""
state = get_judge_task_state(app)
if source_id in state.running_book_ids:
logger.info(f"ebook_book_phrase_judgment_already_running {source_id=}")
return False
state.running_book_ids.add(source_id)
state.outcome_messages.pop(source_id, None)
background_tasks.add_task(judge_book_phrases_for_app, app, source_id)
logger.info(f"ebook_book_phrase_judgment_queued {source_id=}")
return True
async def judge_book_phrases_for_app(app: FastAPI, source_id: int) -> None:
"""Judge one book using the app engine and config, recording the outcome message.
Args:
app (FastAPI): App supplying the engine, config, and judge task state.
source_id (int): Book to judge candidates for.
"""
state = get_judge_task_state(app)
try:
result = await judge_candidate_phrases_for_books(app.state.engine, app.state.config, source_ids=[source_id])
logger.info(
f"ebook_book_phrase_judgment_complete {source_id=} {result.candidates_judged=} {result.protected_phrases=} "
f"{result.phrase_mentions=} {result.books_failed=}"
)
if result.books_failed:
message = "Judging failed; see server logs for details"
else:
message = (
f"Judged {result.candidates_judged} candidates; {result.protected_phrases} protected phrases promoted"
)
except Exception:
logger.exception(f"ebook_book_phrase_judgment_task_failed {source_id=}")
message = "Judging failed; see server logs for details"
state.running_book_ids.discard(source_id)
state.outcome_messages[source_id] = message
def is_judging_book(app: FastAPI, source_id: int) -> bool:
"""Report whether a judgment is currently queued or running for one book.
Args:
app (FastAPI): App supplying the judge task state.
source_id (int): Book to check.
Returns:
bool: True while the book's judgment is pending or running.
"""
return source_id in get_judge_task_state(app).running_book_ids
def pop_book_judgment_outcome(app: FastAPI, source_id: int) -> str | None:
"""Return and clear the outcome message from one book's last finished judgment.
Args:
app (FastAPI): App supplying the judge task state.
source_id (int): Book to fetch the outcome for.
Returns:
str | None: The outcome message, or None when there is nothing new to report.
"""
return get_judge_task_state(app).outcome_messages.pop(source_id, None)