dotfiles/python/ebook_search/api/judge_tasks.py

"""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("ebook_book_phrase_judgment_already_running source_id=%s", 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("ebook_book_phrase_judgment_queued source_id=%s", 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(
            "ebook_book_phrase_judgment_complete source_id=%s judged=%s protected=%s mentions=%s failed=%s",
            source_id,
            result.candidates_judged,
            result.protected_phrases,
            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("ebook_book_phrase_judgment_task_failed source_id=%s", 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)