Add analysis engine: batch pipeline, interactive Q&A, context builder

analysis/context_builder.py

@@ -0,0 +1,168 @@
+"""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,
+        }