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
130 lines
4.4 KiB
Python
130 lines
4.4 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,
|
|
AppEngine,
|
|
AppHttpClient,
|
|
)
|
|
from python.ebook_search.api.web import 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(
|
|
"ebook_answer_low_confidence confidence=%.4f threshold=%.4f",
|
|
retrieval_confidence(response.results),
|
|
config.min_retrieval_confidence,
|
|
)
|
|
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("ebook_answer_request_failed_falling_back error=%s", 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(
|
|
"ebook_answer_citation_issue invalid=%s grounded=%s",
|
|
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[str | None, Form()] = None,
|
|
phrase_matching: Annotated[str | None, Form()] = None,
|
|
) -> HTMLResponse:
|
|
"""Run a search and render HTMX results."""
|
|
try:
|
|
response = await search_ebooks(
|
|
engine,
|
|
client,
|
|
query,
|
|
config,
|
|
rerank=rerank == "true",
|
|
phrase_matching=phrase_matching == "true",
|
|
)
|
|
except Exception as error:
|
|
logger.exception("ebook_search_request_failed")
|
|
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
|
|
|
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("ebook_search_timing step=%r runtime_ms=%.1f", step.name, step.duration_ms)
|
|
logger.info(
|
|
"ebook_search_request_complete results=%s rank_label=%s runtime_ms=%.1f",
|
|
len(response.results),
|
|
response.rank_label,
|
|
response.total_runtime_ms,
|
|
)
|
|
return templates.TemplateResponse(
|
|
request,
|
|
"partials/results.html",
|
|
{
|
|
"answer": answer,
|
|
"response": response,
|
|
"low_confidence": low_confidence,
|
|
"citation_report": citation_report,
|
|
},
|
|
)
|