dotfiles/python/ebook_search/answer.py

"""Grounded answer generation."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from python.ebook_search.llm_interface import request_chat_completion

if TYPE_CHECKING:
    import httpx

    from python.ebook_search.config import EbookSearchConfig
    from python.ebook_search.search import SearchResult

logger = logging.getLogger(__name__)


async def answer_query(
    client: httpx.AsyncClient,
    query: str,
    results: list[SearchResult],
    config: EbookSearchConfig,
) -> str:
    """Answer a question using only retrieved chunks."""
    if not config.answer_enabled:
        logger.info("ebook_answer_skipped_disabled")
        return "Answer generation is disabled. Source chunks are shown below."

    if not results:
        logger.info("ebook_answer_skipped_no_results")
        return "No relevant sources were found."

    logger.info(
        f"ebook_answer_request_start {config.vllm_base_url=} {config.chat_model=} sources={len(results)} "
        f"query_length={len(query)}"
    )
    context = "\n\n".join(
        f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}"
        for index, result in enumerate(results, start=1)
    )
    content = await request_chat_completion(
        client,
        config,
        [
            {
                "role": "system",
                "content": (
                    "Answer only from the provided context. Cite sources with bracketed numbers like [1]. "
                    "If the context is insufficient, say so."
                ),
            },
            {"role": "user", "content": f"Question:\n{query}\n\nContext:\n{context}"},
        ],
    )

    logger.info(f"ebook_answer_request_complete {config.chat_model=} answer_length={len(content)}")
    return content or "The model returned an empty answer."