dotfiles/python/ebook_search/api/routes/search.py

"""Search routes for the EPUB search web UI."""

from __future__ import annotations

import logging
from dataclasses import replace
from time import perf_counter
from typing import TYPE_CHECKING, Annotated

from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse

from python.ebook_search.answer import answer_query
from python.ebook_search.api.dependencies import (  # noqa: TC001 FastAPI resolves these annotated dependencies at runtime
    AppConfig,
    AppEngine,
    AppHttpClient,
)
from python.ebook_search.api.web import error_response, templates
from python.ebook_search.guardrails import (
    CitationReport,
    is_confident,
    retrieval_confidence,
    validate_citations,
)
from python.ebook_search.search import SearchResponse, search_ebooks
from python.ebook_search.timing import runtime_step_from_start

if TYPE_CHECKING:
    import httpx

    from python.ebook_search.config import EbookSearchConfig

logger = logging.getLogger(__name__)

router = APIRouter()


async def build_answer(
    client: httpx.AsyncClient,
    query: str,
    response: SearchResponse,
    config: EbookSearchConfig,
) -> tuple[str, bool, CitationReport | None]:
    """Generate the answer for a search, returning ``(answer, low_confidence, citation_report)``."""
    if not config.answer_enabled:
        logger.info("ebook_answer_skipped_disabled")
        return "Answer generation is disabled. Source chunks are shown below.", False, None

    if not is_confident(response.results, config):
        logger.info(
            f"ebook_answer_low_confidence confidence={retrieval_confidence(response.results):.4f} "
            f"{config.min_retrieval_confidence=:.4f}"
        )
        answer = (
            "Retrieval confidence is low for this query, so answer generation was skipped. "
            "Source chunks are shown below."
        )
        return answer, True, None

    try:
        answer = await answer_query(client, query, response.results, config)
    except RuntimeError as error:
        logger.warning(f"ebook_answer_request_failed_falling_back {error=}")
        return "Answer generation failed. Source chunks are still shown below.", False, None

    citation_report = None
    if config.validate_citations_enabled and response.results:
        citation_report = validate_citations(answer, len(response.results))
        if citation_report.invalid or not citation_report.grounded:
            logger.warning(f"ebook_answer_citation_issue {citation_report.invalid=} {citation_report.grounded=}")
    return answer, False, citation_report


@router.post("/search", response_class=HTMLResponse)
async def search(
    request: Request,
    config: AppConfig,
    engine: AppEngine,
    client: AppHttpClient,
    query: Annotated[str, Form()],
    *,
    rerank: Annotated[bool, Form()] = False,
    phrase_matching: Annotated[bool, Form()] = False,
) -> HTMLResponse:
    """Run a search and render HTMX results."""
    try:
        response = await search_ebooks(
            engine,
            client,
            query,
            config,
            rerank=rerank,
            phrase_matching=phrase_matching,
        )
    except Exception as error:
        logger.exception("ebook_search_request_failed")
        return error_response(request, error)

    answer_start = perf_counter()
    answer, low_confidence, citation_report = await build_answer(client, query, response, config)
    answer_step_name = "Answer generation" if config.answer_enabled else "Answer skipped"
    response = replace(
        response,
        timings=(*response.timings, runtime_step_from_start(answer_step_name, answer_start)),
    )

    for step in response.timings:
        logger.info(f"ebook_search_timing {step.name=} {step.duration_ms=:.1f}")
    logger.info(
        f"ebook_search_request_complete results={len(response.results)} {response.rank_label=} "
        f"{response.total_runtime_ms=:.1f}"
    )
    return templates.TemplateResponse(
        request,
        "partials/results.html",
        {
            "answer": answer,
            "response": response,
            "low_confidence": low_confidence,
            "citation_report": citation_report,
        },
    )