mosaic/analysis/context_builder.py

"""Context builder for assembling relevant document chunks for Claude Q&A."""
import logging
from typing import Optional

from storage.db import MosaicDB
from utils.text_utils import estimate_tokens

logger = logging.getLogger(__name__)

DEFAULT_TOKEN_BUDGET = 30000
DEFAULT_MAX_RESULTS = 20


class ContextBuilder:
    """Build context from FTS5 search results for Claude queries.

    Uses BM25 ranking from FTS5, assembles chunks within a token budget,
    and tracks source citations.
    """

    def __init__(self, db: MosaicDB, token_budget: int = DEFAULT_TOKEN_BUDGET,
                 max_results: int = DEFAULT_MAX_RESULTS):
        self.db = db
        self.token_budget = token_budget
        self.max_results = max_results

    def build_context(self, query: str, source: Optional[str] = None) -> dict:
        """Search documents and assemble context within token budget.

        Args:
            query: Search query string
            source: Optional source filter

        Returns:
            dict with:
                'context_text': assembled context string with citations
                'documents': list of matched document metadata
                'total_tokens': estimated token count
                'truncated': whether results were truncated
        """
        # Search with FTS5
        docs = self.db.search_documents(query, source=source, limit=self.max_results)

        if not docs:
            return {
                'context_text': '',
                'documents': [],
                'total_tokens': 0,
                'truncated': False,
            }

        # Assemble context within budget
        context_parts = []
        used_tokens = 0
        included_docs = []
        truncated = False

        for doc in docs:
            # Build document excerpt
            excerpt = self._format_document_excerpt(doc)
            excerpt_tokens = estimate_tokens(excerpt)

            if used_tokens + excerpt_tokens > self.token_budget:
                # Try to fit a truncated version
                remaining_tokens = self.token_budget - used_tokens
                if remaining_tokens > 500:  # Worth including a partial
                    truncated_excerpt = self._truncate_to_tokens(excerpt, remaining_tokens)
                    context_parts.append(truncated_excerpt)
                    used_tokens += remaining_tokens
                    included_docs.append(self._doc_metadata(doc))
                truncated = True
                break

            context_parts.append(excerpt)
            used_tokens += excerpt_tokens
            included_docs.append(self._doc_metadata(doc))

        context_text = '\n\n---\n\n'.join(context_parts)

        return {
            'context_text': context_text,
            'documents': included_docs,
            'total_tokens': used_tokens,
            'truncated': truncated,
        }

    def build_entity_context(self, entity_name: str) -> str:
        """Build context focused on a specific entity across all sources."""
        docs = self.db.search_documents(entity_name, limit=self.max_results)

        parts = []
        used_tokens = 0
        for doc in docs:
            excerpt = self._format_document_excerpt(doc)
            tokens = estimate_tokens(excerpt)
            if used_tokens + tokens > self.token_budget:
                break
            parts.append(excerpt)
            used_tokens += tokens

        return '\n\n---\n\n'.join(parts)

    def build_tool_context(self, tool_name: str) -> str:
        """Build context focused on a specific tool/implant."""
        # Search both FTS and extracted tools
        docs = self.db.search_documents(tool_name, limit=self.max_results)

        parts = []
        used_tokens = 0

        # Add extracted tool data if available
        tools = self.db.get_all_tools()
        for tool in tools:
            if tool_name.lower() in tool.name.lower() or \
               any(tool_name.lower() in a.lower() for a in tool.aliases):
                tool_text = f"[Extracted Tool: {tool.name}]\n"
                tool_text += f"Capability: {tool.capability}\n"
                tool_text += f"Description: {tool.description}\n"
                tool_text += f"Platforms: {', '.join(tool.target_platforms)}\n"
                tool_text += f"CVEs: {', '.join(tool.cves)}\n"
                tool_text += f"Source context: {tool.source_context}\n"
                parts.append(tool_text)
                used_tokens += estimate_tokens(tool_text)

        # Add document excerpts
        for doc in docs:
            excerpt = self._format_document_excerpt(doc)
            tokens = estimate_tokens(excerpt)
            if used_tokens + tokens > self.token_budget:
                break
            parts.append(excerpt)
            used_tokens += tokens

        return '\n\n---\n\n'.join(parts)

    def _format_document_excerpt(self, doc) -> str:
        """Format a document for context inclusion with citation info."""
        header = f"[Source: {doc.source} | Title: {doc.title}"
        if doc.date:
            header += f" | Date: {doc.date}"
        if doc.classification:
            header += f" | Classification: {doc.classification}"
        header += f" | ID: {doc.id[:12]}]"

        # Truncate text if needed
        text = doc.text
        max_chars = (self.token_budget // len(self.db.search_documents("", limit=1) or [doc])) * 4
        if len(text) > max_chars:
            text = text[:max_chars] + "\n[...truncated...]"

        return f"{header}\n\n{text}"

    def _truncate_to_tokens(self, text: str, max_tokens: int) -> str:
        """Truncate text to approximately max_tokens."""
        max_chars = max_tokens * 4  # Rough chars-per-token estimate
        if len(text) <= max_chars:
            return text
        return text[:max_chars] + "\n[...truncated...]"

    def _doc_metadata(self, doc) -> dict:
        """Extract metadata dict from a document for citation tracking."""
        return {
            'id': doc.id,
            'title': doc.title,
            'source': doc.source,
            'date': doc.date,
            'doc_type': doc.doc_type,
        }