generation/generator.py

import logging
from typing import Any, Dict, List

from .llm_client import LLMClient, LLMClientError
from .prompt_builder import build_prompt, format_context, MAX_CHUNKS

logger = logging.getLogger(__name__)


class GenerationError(Exception):
    pass


class Generator:
    def __init__(self, llm_client: LLMClient = None) -> None:
        self.llm_client = llm_client or LLMClient()
        logger.info("Generator ready.")

    def generate(
        self,
        query: str,
        chunks: List[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Run the full generation step: prompt → LLM → structured result.

        Returns:
            {
                "query":    original question,
                "answer":   model's answer string,
                "contexts": list of chunk_text strings sent to the model,
                "sources":  list of source filenames sent to the model,
            }
        """
        if not query or not query.strip():
            raise ValueError("query must be a non-empty string.")

        # Mirror format_context's selection so contexts/sources match the prompt
        ordered = sorted(chunks, key=lambda c: c.get("rank", 0))[:MAX_CHUNKS]
        used = [c for c in ordered if (c.get("chunk_text") or "").strip()]

        messages = build_prompt(query, chunks)

        try:
            answer = self.llm_client.generate(messages)
        except LLMClientError as exc:
            logger.error("Generation failed for query %r: %s", query, exc)
            raise GenerationError(f"Failed to generate answer: {exc}") from exc

        logger.info("Answer generated (%d chars, %d chunks used).",
                    len(answer), len(used))

        return {
            "query": query.strip(),
            "answer": answer,
            "contexts": [c.get("chunk_text", "") for c in used],
            "sources": [c.get("source", "unknown") for c in used],
        }