"""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()