core/context_builder.py

import logging
from typing import List

logger = logging.getLogger("axiom.context")


class ContextBuilder:
    """
    Assembles retrieved chunks into a structured context string
    for the LLM, with document source tracking.
    """

    MAX_CONTEXT_WORDS = 1500

    def __init__(self):
        logger.info("Ready.")

    def _build_context(self, chunks: List[dict]) -> tuple[str, List[str]]:
        context_parts = []
        sources       = []
        word_count    = 0

        for chunk in chunks:
            chunk_words = chunk["text"].split()

            if word_count + len(chunk_words) > self.MAX_CONTEXT_WORDS:
                remaining = self.MAX_CONTEXT_WORDS - word_count
                if remaining > 20:
                    truncated = " ".join(chunk_words[:remaining])
                    context_parts.append(truncated)
                title = chunk["metadata"].get("title", "Unknown")
                if title not in sources:
                    sources.append(title)
                break

            context_parts.append(chunk["text"])
            word_count += len(chunk_words)

            title = chunk["metadata"].get("title", "Unknown")
            if title not in sources:
                sources.append(title)

        context_text = " ".join(context_parts)
        return context_text, sources

    def build(self, query: str, chunks: List[dict]) -> dict:
        context_text, sources = self._build_context(chunks)

        prompt = f"""Use the context below to answer the question accurately.

Context: {context_text}

Question: {query}

Provide a factual answer based strictly on the context above:"""

        return {
            "prompt":       prompt,
            "context_text": context_text,
            "sources":      sources,
            "word_count":   len(context_text.split())
        }