n0mad1k
259 lines
10 KiB
Python
259 lines
10 KiB
Python
"""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()
|