context_manager.py

from typing import List, Dict, Any, Optional


def _estimate_tokens(text: str) -> int:
    return len(text) // 4 + 1


class ContextManager:
    def __init__(self, max_tokens: int = 8192):
        self.max_tokens = max_tokens

    def assemble_prompt(
        self,
        query: str,
        documents: List[Dict[str, Any]],
        system_prompt: Optional[str] = None,
    ) -> str:
        if system_prompt is None:
            system_prompt = (
                "You are a helpful AI assistant. Answer the user's question "
                "based solely on the provided context. Keep your answer concise "
                "and directly address the question. If the context lacks "
                "sufficient information, state that clearly."
            )

        system_tokens = _estimate_tokens(system_prompt)
        query_text = f"Question: {query}\n\nAnswer:"
        query_tokens = _estimate_tokens(query_text)

        budget = self.max_tokens - system_tokens - query_tokens - 50
        if budget < 100:
            budget = 100

        context_parts = []
        chars_used = 0
        budget_chars = budget * 4

        for i, doc in enumerate(documents):
            raw_path = doc.get("metadata", {}).get("hierarchy_path", "")
            path_parts = [p.strip() for p in raw_path.split("|") if p.strip()] if raw_path else []
            path_str = " > ".join(path_parts) if path_parts else ""
            header = f"[Document {i + 1}]" + (f" — {path_str}" if path_str else "")
            header_chars = len(header) + 1
            text = doc.get("text", "")
            text_chars = len(text)
            total_needed = header_chars + text_chars

            if chars_used + total_needed > budget_chars:
                remaining = budget_chars - chars_used - header_chars
                if remaining > 80 and len(context_parts) > 0:
                    truncated = text[:remaining]
                    context_parts.append(f"{header}\n{truncated}")
                    chars_used += header_chars + remaining
                break

            context_parts.append(f"{header}\n{text}")
            chars_used += total_needed

        context = "\n\n".join(context_parts)
        return f"{system_prompt}\n\nContext:\n{context}\n\n{query_text}"

    def count_tokens(self, text: str) -> int:
        return _estimate_tokens(text)