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

analysis/__init__.py

@@ -0,0 +1,3 @@
+from .claude_batch import BatchAnalyzer
+from .claude_interactive import InteractiveAnalyzer
+from .context_builder import ContextBuilder

analysis/claude_batch.py

@@ -0,0 +1,258 @@
+"""Batch document analysis pipeline for Mosaic."""
+import logging
+from typing import Optional
+from datetime import datetime
+
+from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
+from rich.console import Console
+
+from storage.db import MosaicDB
+from storage.models import Document
+from storage.jsonl_store import JSONLStore
+from utils.api_limiter import APIRateLimiter, CircuitBreakerOpen
+from utils.text_utils import estimate_tokens
+
+logger = logging.getLogger(__name__)
+console = Console()
+
+EXTRACTOR_DOMAINS = {
+    'tools': 'tool',
+    'ttps': 'ttp',
+    'tradecraft': 'tradecraft',
+    'surveillance': 'surveillance',
+    'infrastructure': 'infrastructure',
+    'entities': 'entity',
+}
+
+
+class BatchAnalyzer:
+    """Run batch extraction pipeline across documents."""
+
+    def __init__(self, db: MosaicDB, api_key: str,
+                 model: str = "claude-sonnet-4-20250514",
+                 rate_limiter: Optional[APIRateLimiter] = None):
+        self.db = db
+        self.api_key = api_key
+        self.model = model
+        self.rate_limiter = rate_limiter or APIRateLimiter()
+        self._extractors = {}
+
+    def _get_extractor(self, domain: str):
+        """Lazy-load extractor by domain name."""
+        if domain not in self._extractors:
+            from extractors.ttp_extractor import TTPExtractor
+            from extractors.tool_extractor import ToolExtractor
+            from extractors.tradecraft import TradecraftExtractor
+            from extractors.surveillance import SurveillanceExtractor
+            from extractors.infrastructure import InfrastructureExtractor
+            from extractors.entity_extractor import EntityExtractor
+
+            extractor_classes = {
+                'ttp': TTPExtractor,
+                'tool': ToolExtractor,
+                'tradecraft': TradecraftExtractor,
+                'surveillance': SurveillanceExtractor,
+                'infrastructure': InfrastructureExtractor,
+                'entity': EntityExtractor,
+            }
+
+            cls = extractor_classes.get(domain)
+            if cls:
+                if domain == 'infrastructure':
+                    self._extractors[domain] = cls(
+                        db=self.db, api_key=self.api_key,
+                        model=self.model, rate_limiter=self.rate_limiter,
+                        collection="", doc_title="", doc_date=None,
+                    )
+                else:
+                    self._extractors[domain] = cls(
+                        db=self.db, api_key=self.api_key,
+                        model=self.model, rate_limiter=self.rate_limiter,
+                    )
+        return self._extractors.get(domain)
+
+    def estimate(self, source: Optional[str] = None,
+                 domains: Optional[list[str]] = None,
+                 since: Optional[str] = None,
+                 until: Optional[str] = None,
+                 since_last_run: bool = False) -> dict:
+        """Estimate token costs without running extraction (dry-run).
+
+        Returns dict with document count, chunk count, estimated tokens.
+        """
+        domains = domains or list(EXTRACTOR_DOMAINS.values())
+        docs = self._get_target_documents(source, since, until, since_last_run)
+
+        total_chunks = 0
+        total_input_tokens = 0
+
+        for doc in docs:
+            for domain in domains:
+                extractor = self._get_extractor(domain)
+                if extractor:
+                    est = extractor.estimate_tokens(doc.text, doc.doc_type)
+                    total_chunks += est['chunks']
+                    total_input_tokens += est['estimated_input_tokens']
+
+        return {
+            'document_count': len(docs),
+            'domain_count': len(domains),
+            'total_chunks': total_chunks,
+            'estimated_input_tokens': total_input_tokens,
+            'estimated_output_tokens': total_chunks * 2000,
+            'estimated_total_tokens': total_input_tokens + (total_chunks * 2000),
+        }
+
+    def run(self, source: Optional[str] = None,
+            domains: Optional[list[str]] = None,
+            since: Optional[str] = None,
+            until: Optional[str] = None,
+            since_last_run: bool = False,
+            retry_failed: bool = False) -> dict:
+        """Run batch extraction pipeline.
+
+        Returns dict with counts of extracted items per domain.
+        """
+        domains = domains or list(EXTRACTOR_DOMAINS.values())
+        docs = self._get_target_documents(source, since, until, since_last_run)
+
+        if not docs:
+            console.print("[yellow]No documents to process[/yellow]")
+            return {}
+
+        results = {d: 0 for d in domains}
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+        ) as progress:
+            for domain in domains:
+                extractor = self._get_extractor(domain)
+                if not extractor:
+                    console.print(f"[yellow]Unknown domain: {domain}[/yellow]")
+                    continue
+
+                task = progress.add_task(f"Extracting {domain}", total=len(docs))
+
+                for doc in docs:
+                    # Check if already done
+                    status = self.db.get_extraction_status(doc.id, extractor.EXTRACTOR_NAME)
+                    if status and status.status == 'DONE' and not retry_failed:
+                        progress.advance(task)
+                        continue
+                    if status and status.status == 'FAILED' and not retry_failed:
+                        progress.advance(task)
+                        continue
+
+                    try:
+                        # Update infrastructure extractor context
+                        if domain == 'infrastructure' and hasattr(extractor, 'collection'):
+                            extractor.collection = doc.source
+                            extractor.doc_title = doc.title
+                            extractor.doc_date = doc.date
+
+                        count = extractor.extract_from_document(
+                            doc_id=doc.id,
+                            text=doc.text,
+                            doc_type=doc.doc_type,
+                            source=doc.source,
+                            metadata=doc.metadata,
+                        )
+                        results[domain] += count
+
+                    except CircuitBreakerOpen:
+                        console.print(f"[red]Circuit breaker open — stopping {domain} extraction[/red]")
+                        break
+                    except Exception as e:
+                        logger.error("Extraction error for %s on doc %s: %s",
+                                   domain, doc.id[:12], e)
+
+                    progress.advance(task)
+
+        return results
+
+    def run_full_pipeline(self, source: str, limit: Optional[int] = None) -> dict:
+        """Run full pipeline: collect → parse → extract → report.
+
+        Returns summary dict.
+        """
+        from collectors import load_profile, CustomSource
+        from parsers import ParserManager
+        from utils.cache import DownloadCache
+        from utils.http_limiter import HTTPRateLimiter
+
+        summary = {'collect': {}, 'parse': {}, 'extract': {}, 'export': []}
+
+        # Step 1: Collect
+        console.print(f"\n[bold cyan]Step 1/4: Collecting from {source}[/bold cyan]")
+        try:
+            profile = load_profile(source)
+            cache = DownloadCache()
+            collector = CustomSource(profile, self.db, cache)
+            summary['collect'] = collector.collect(limit=limit)
+        except Exception as e:
+            console.print(f"[red]Collection error: {e}[/red]")
+            return summary
+
+        # Step 2: Parse
+        console.print(f"\n[bold cyan]Step 2/4: Parsing documents[/bold cyan]")
+        parser = ParserManager()
+        docs = self.db.get_documents_by_source(source, status='CACHED')
+        parsed = 0
+        failed = 0
+
+        with Progress(SpinnerColumn(), TextColumn("{task.description}"),
+                      BarColumn(), TaskProgressColumn()) as progress:
+            task = progress.add_task("Parsing", total=len(docs))
+            for doc in docs:
+                if doc.parse_status in ('PARSED', 'NEEDS_OCR'):
+                    progress.advance(task)
+                    continue
+
+                self.db.update_document_status(doc.id, parse_status='IN_PROGRESS')
+                result = parser.parse(doc.raw_path, doc.doc_type, doc.content_hash)
+
+                if result.success:
+                    doc.text = result.text
+                    doc.title = result.title or doc.title
+                    doc.char_count = result.char_count
+                    doc.metadata.update(result.metadata)
+                    status = 'NEEDS_OCR' if result.needs_ocr else 'PARSED'
+                    doc.parse_status = status
+                    self.db.insert_document(doc)
+                    self.db.update_document_status(doc.id, parse_status=status)
+                    parsed += 1
+                else:
+                    self.db.update_document_status(doc.id, parse_status='PARSE_FAILED')
+                    failed += 1
+
+                progress.advance(task)
+
+        summary['parse'] = {'parsed': parsed, 'failed': failed}
+
+        # Step 3: Extract
+        console.print(f"\n[bold cyan]Step 3/4: Running extraction[/bold cyan]")
+        summary['extract'] = self.run(source=source)
+
+        # Step 4: Export
+        console.print(f"\n[bold cyan]Step 4/4: Exporting results[/bold cyan]")
+        store = JSONLStore(self.db)
+        summary['export'] = store.export()
+
+        return summary
+
+    def _get_target_documents(self, source: Optional[str] = None,
+                               since: Optional[str] = None,
+                               until: Optional[str] = None,
+                               since_last_run: bool = False) -> list[Document]:
+        """Get documents to process based on filters."""
+        if since or until:
+            return self.db.get_documents_by_date_range(source=source, since=since, until=until)
+        elif since_last_run:
+            return self.db.get_unextracted_documents(source=source, since_last_run=True)
+        elif source:
+            return self.db.get_documents_by_source(source, parse_status='PARSED')
+        else:
+            return self.db.get_unextracted_documents()

analysis/claude_interactive.py

@@ -0,0 +1,206 @@
+"""Interactive Q&A mode for Mosaic — Rich-based chat over extracted intelligence."""
+import logging
+from typing import Optional
+
+import anthropic
+from rich.console import Console
+from rich.panel import Panel
+from rich.markdown import Markdown
+from rich.prompt import Prompt
+
+from storage.db import MosaicDB
+from .context_builder import ContextBuilder
+
+logger = logging.getLogger(__name__)
+console = Console()
+
+# Import prompts — handle if not yet created
+try:
+    from .prompts import (
+        QUERY_SYSTEM_PROMPT, QUERY_PROMPT, EMULATE_PROMPT,
+        COMPARE_PROMPT, BRIEF_PROMPT, SCENARIO_PROMPT
+    )
+except ImportError:
+    QUERY_SYSTEM_PROMPT = "You are an intelligence analyst. Answer questions based on provided documents."
+    QUERY_PROMPT = "{context}\n\nQuestion: {question}"
+    EMULATE_PROMPT = "Emulate threat actor {actor}:\n{context}"
+    COMPARE_PROMPT = "Compare {tool_a} vs {tool_b}:\n{context}"
+    BRIEF_PROMPT = "Brief on {topic}:\n{context}"
+    SCENARIO_PROMPT = "Scenario for {target}:\n{context}"
+
+
+class InteractiveAnalyzer:
+    """Interactive Q&A over extracted intelligence corpus."""
+
+    def __init__(self, db: MosaicDB, api_key: str,
+                 model: str = "claude-sonnet-4-20250514",
+                 source: Optional[str] = None):
+        self.db = db
+        self.client = anthropic.Anthropic(api_key=api_key)
+        self.model = model
+        self.source = source
+        self.context_builder = ContextBuilder(db)
+        self.mode = "default"
+        self.conversation_history = []
+
+    def run(self):
+        """Start interactive Q&A session."""
+        self._show_welcome()
+
+        while True:
+            try:
+                user_input = Prompt.ask("\n[bold green]mosaic[/bold green]")
+            except (KeyboardInterrupt, EOFError):
+                console.print("\n[dim]Goodbye.[/dim]")
+                break
+
+            user_input = user_input.strip()
+            if not user_input:
+                continue
+
+            # Handle commands
+            if user_input.lower() in ('exit', 'quit', 'q'):
+                console.print("[dim]Goodbye.[/dim]")
+                break
+            elif user_input.lower() == 'help':
+                self._show_help()
+                continue
+            elif user_input.lower().startswith('mode '):
+                self.mode = user_input[5:].strip()
+                console.print(f"[cyan]Mode set to: {self.mode}[/cyan]")
+                continue
+            elif user_input.lower().startswith('source '):
+                self.source = user_input[7:].strip()
+                console.print(f"[cyan]Source filter: {self.source}[/cyan]")
+                continue
+            elif user_input.lower() == 'clear':
+                self.conversation_history = []
+                console.print("[dim]Conversation cleared.[/dim]")
+                continue
+            elif user_input.lower() == 'status':
+                self._show_status()
+                continue
+
+            # Process query based on mode
+            self._handle_query(user_input)
+
+    def _handle_query(self, query: str):
+        """Process a query in the current mode."""
+        with console.status("[bold cyan]Searching corpus...[/bold cyan]"):
+            # Build context based on mode
+            if self.mode == 'emulate':
+                context = self.context_builder.build_context(query, self.source)
+                prompt = EMULATE_PROMPT.format(actor=query, context=context['context_text'])
+            elif self.mode == 'compare' and ' vs ' in query:
+                parts = query.split(' vs ', 1)
+                tool_a = parts[0].strip()
+                tool_b = parts[1].strip()
+                ctx_a = self.context_builder.build_tool_context(tool_a)
+                ctx_b = self.context_builder.build_tool_context(tool_b)
+                context = {'context_text': ctx_a + '\n\n' + ctx_b, 'documents': [], 'total_tokens': 0}
+                prompt = COMPARE_PROMPT.format(tool_a=tool_a, tool_b=tool_b, context=context['context_text'])
+            elif self.mode == 'brief':
+                context = self.context_builder.build_context(query, self.source)
+                prompt = BRIEF_PROMPT.format(topic=query, context=context['context_text'])
+            elif self.mode == 'scenario':
+                context = self.context_builder.build_context(query, self.source)
+                prompt = SCENARIO_PROMPT.format(target=query, context=context['context_text'])
+            else:
+                context = self.context_builder.build_context(query, self.source)
+                prompt = QUERY_PROMPT.format(context=context['context_text'], question=query)
+
+        if not context.get('context_text'):
+            console.print("[yellow]No relevant documents found for your query.[/yellow]")
+            return
+
+        # Show source count
+        doc_count = len(context.get('documents', []))
+        console.print(f"[dim]Found {doc_count} relevant documents, "
+                     f"~{context.get('total_tokens', 0):,} tokens of context[/dim]")
+
+        # Call Claude
+        with console.status("[bold cyan]Analyzing...[/bold cyan]"):
+            try:
+                messages = list(self.conversation_history)
+                messages.append({"role": "user", "content": prompt})
+
+                response = self.client.messages.create(
+                    model=self.model,
+                    max_tokens=4096,
+                    system=QUERY_SYSTEM_PROMPT,
+                    messages=messages,
+                )
+
+                answer = response.content[0].text
+
+                # Save to conversation history
+                self.conversation_history.append({"role": "user", "content": prompt})
+                self.conversation_history.append({"role": "assistant", "content": answer})
+
+                # Trim history if too long
+                if len(self.conversation_history) > 20:
+                    self.conversation_history = self.conversation_history[-10:]
+
+            except anthropic.APIError as e:
+                console.print(f"[red]API error: {e}[/red]")
+                return
+
+        # Display response
+        console.print()
+        console.print(Panel(
+            Markdown(answer),
+            title="[bold]Analysis[/bold]",
+            border_style="cyan",
+        ))
+
+        # Show citations
+        if context.get('documents'):
+            console.print("\n[dim]Sources:[/dim]")
+            for doc in context['documents'][:5]:
+                console.print(f"  [dim]- {doc['source']}: {doc['title']}", markup=False)
+
+    def _show_welcome(self):
+        """Show welcome message."""
+        welcome = """[bold cyan]MOSAIC Interactive Intelligence Query[/bold cyan]
+
+Ask questions about the extracted intelligence corpus.
+Type [bold]help[/bold] for commands, [bold]exit[/bold] to quit."""
+
+        if self.source:
+            welcome += f"\n[dim]Source filter: {self.source}[/dim]"
+
+        console.print(Panel(welcome, border_style="cyan"))
+
+    def _show_help(self):
+        """Show help text."""
+        help_text = """[bold]Commands:[/bold]
+  help              Show this help
+  exit/quit         Leave interactive mode
+  clear             Clear conversation history
+  status            Show corpus statistics
+  mode <mode>       Set query mode:
+                      default   - General Q&A
+                      emulate   - Threat actor emulation playbook
+                      compare   - Compare tools (use 'tool_a vs tool_b')
+                      brief     - Intelligence briefing
+                      scenario  - Adversary simulation scenario
+  source <name>     Filter to specific source (or 'all')
+
+[bold]Tips:[/bold]
+  - Ask specific questions for better results
+  - Use source names (vault7, cablegate) to focus queries
+  - Emulate mode: type an actor name to get their TTP playbook"""
+        console.print(Panel(help_text, border_style="cyan"))
+
+    def _show_status(self):
+        """Show corpus statistics."""
+        stats = self.db.get_status_counts(self.source)
+        tokens = self.db.get_total_tokens(self.source)
+
+        console.print(f"\n[bold]Corpus Statistics[/bold]")
+        console.print(f"  Documents: {stats.get('total', 0)}")
+        console.print(f"  Parsed: {stats.get('parsing', {}).get('PARSED', 0)}")
+        console.print(f"  Tokens used: {tokens.get('total_tokens', 0):,}")
+
+        if self.source:
+            console.print(f"  [dim]Filtered to: {self.source}[/dim]")

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,
+        }