Add analysis engine: batch pipeline, interactive Q&A, context builder
This commit is contained in:
n0mad1k
@@ -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()
|
||||
Reference in New Issue
Block a user