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

124 lines
4.2 KiB
Python

"""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,
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
from python.fastapi_tools import AppAsyncEngine # noqa: TC001 FastAPI resolves this annotated dependency at runtime
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: AppAsyncEngine,
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,
},
)