added guardrails.py to constrain responses and added validation to config.py

python/ebook_search/api/routes/search.py

@@ -11,8 +11,16 @@ 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 AppConfig, AppEngine
 from python.ebook_search.api.web import templates
-from python.ebook_search.search import search_ebooks
+from python.ebook_search.config import EbookSearchConfig
+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
 
 logger = logging.getLogger(__name__)
@@ -20,30 +28,64 @@ logger = logging.getLogger(__name__)
 router = APIRouter()
 
 
+def build_answer(
+    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 = answer_query(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)
 def search(
     request: Request,
+    config: AppConfig,
+    engine: AppEngine,
     query: Annotated[str, Form()],
     rerank: Annotated[str | None, Form()] = None,
 ) -> HTMLResponse:
     """Run a search and render HTMX results."""
     try:
-        response = search_ebooks(request.app.state.engine, query, request.app.state.config, rerank=rerank == "true")
+        response = search_ebooks(engine, query, config, rerank=rerank == "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()
-    if request.app.state.config.answer_enabled:
-        try:
-            answer = answer_query(query, response.results, request.app.state.config)
-        except RuntimeError as error:
-            logger.warning("ebook_answer_request_failed_falling_back error=%s", error)
-            answer = "Answer generation failed. Source chunks are still shown below."
-    else:
-        logger.info("ebook_answer_skipped_disabled")
-        answer = "Answer generation is disabled. Source chunks are shown below."
-    answer_step_name = "Answer generation" if request.app.state.config.answer_enabled else "Answer skipped"
+    answer, low_confidence, citation_report = build_answer(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)),
@@ -55,4 +97,13 @@ def search(
         response.rank_label,
         response.total_runtime_ms,
     )
-    return templates.TemplateResponse(request, "partials/results.html", {"answer": answer, "response": response})
+    return templates.TemplateResponse(
+        request,
+        "partials/results.html",
+        {
+            "answer": answer,
+            "response": response,
+            "low_confidence": low_confidence,
+            "citation_report": citation_report,
+        },
+    )