app/services/context_builder.py

"""Author RAG Chatbot SaaS — Token-Aware Context Builder.

Assembles retrieved chunks into a formatted context string that fits
within the configured token budget.
RULE: Hard limit — never exceed RAG_MAX_CONTEXT_TOKENS.
RULE: Include book title for each chunk to help model cross-book navigation.
"""

import structlog

from app.config import get_settings
from app.services.vector_store import RetrievedChunk
from app.utils.token_counter import count_tokens

logger = structlog.get_logger(__name__)
cfg = get_settings()


def build_context(
    chunks: list[RetrievedChunk],
    max_tokens: int | None = None,
) -> tuple[str, int]:
    """Build a formatted context string from retrieved chunks.

    Includes as many chunks as fit within the token budget (best first).
    Each chunk is formatted with a book title header for clarity.

    Args:
        chunks: Re-ranked list of RetrievedChunk objects (best first).
        max_tokens: Max tokens for the context block.

    Returns:
        Tuple of (context_string, total_tokens_used).
    """
    max_tokens = max_tokens or cfg.RAG_MAX_CONTEXT_TOKENS
    included_chunks: list[str] = []
    total_tokens = 0

    for chunk in chunks:
        formatted = _format_chunk(chunk)
        chunk_tokens = count_tokens(formatted)

        if total_tokens + chunk_tokens > max_tokens:
            logger.debug(
                "Context token budget reached",
                included=len(included_chunks),
                excluded_remaining=len(chunks) - len(included_chunks),
            )
            break

        included_chunks.append(formatted)
        total_tokens += chunk_tokens

    if not included_chunks:
        return "", 0

    context = "\n\n---\n\n".join(included_chunks)
    logger.debug("Context built", chunks=len(included_chunks), tokens=total_tokens)
    return context, total_tokens


def _format_chunk(chunk: RetrievedChunk) -> str:
    """Format a single chunk with its book title header.

    Args:
        chunk: RetrievedChunk to format.

    Returns:
        Formatted string with book title and chunk text.
    """
    return f"[From: {chunk.book_title}]\n{chunk.text.strip()}"