#!/usr/bin/env python3 """Mosaic — Intelligence Extraction Platform CLI.""" import sys import os import logging import hashlib import configparser from pathlib import Path from datetime import datetime from typing import Optional import click from rich.console import Console from rich.table import Table from rich.panel import Panel from rich.prompt import Prompt from rich.markdown import Markdown from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn from rich import box # --------------------------------------------------------------------------- # Paths & constants # --------------------------------------------------------------------------- BASE_DIR = Path(__file__).resolve().parent CONFIG_PATH = BASE_DIR / "config" / "mosaic.conf" PROFILES_DIR = BASE_DIR / "collectors" / "profiles" DEFAULT_DB = "mosaic.db" console = Console() logger = logging.getLogger("mosaic") DOMAIN_CHOICES = ["tools", "ttps", "tradecraft", "surveillance", "infrastructure", "entities"] EXPORT_CATEGORIES = ["tools", "ttps", "tradecraft", "entities", "infrastructure"] QUERY_MODES = ["default", "emulate", "compare", "brief", "scenario"] # --------------------------------------------------------------------------- # Config helpers # --------------------------------------------------------------------------- def load_config() -> configparser.ConfigParser: """Load mosaic.conf; return a ConfigParser (empty sections on missing file).""" cfg = configparser.ConfigParser() if CONFIG_PATH.exists(): cfg.read(str(CONFIG_PATH)) else: logger.debug("Config file not found at %s — using defaults", CONFIG_PATH) return cfg def get_api_key(cfg: configparser.ConfigParser) -> Optional[str]: """ANTHROPIC_API_KEY env var takes priority over config file.""" key = os.environ.get("ANTHROPIC_API_KEY") if key: return key return cfg.get("api", "api_key", fallback=None) def get_db_path(cfg: configparser.ConfigParser, cli_override: Optional[str] = None) -> str: """Resolve database path: CLI flag > config > default.""" if cli_override: return cli_override rel = cfg.get("general", "database_path", fallback=DEFAULT_DB) return str(BASE_DIR / rel) def setup_logging(verbose: bool = False): """Configure root logger.""" level = logging.DEBUG if verbose else logging.INFO fmt = "%(asctime)s %(name)-18s %(levelname)-7s %(message)s" logging.basicConfig(level=level, format=fmt, datefmt="%H:%M:%S") def get_output_dir(cfg: configparser.ConfigParser) -> str: """Resolve the output directory from config.""" rel = cfg.get("general", "output_dir", fallback="output") return str(BASE_DIR / rel) # --------------------------------------------------------------------------- # Lazy DB / store helpers (avoid opening DB at import time) # --------------------------------------------------------------------------- def open_db(db_path: str): """Import and return a MosaicDB instance.""" from storage.db import MosaicDB return MosaicDB(db_path) def open_jsonl_store(db, output_dir: str): """Import and return a JSONLStore instance.""" from storage.jsonl_store import JSONLStore return JSONLStore(db, output_dir=os.path.join(output_dir, "extracted")) def open_cache(output_dir: str): """Import and return a DownloadCache instance.""" from utils.cache import DownloadCache return DownloadCache(base_dir=os.path.join(output_dir, "raw")) # --------------------------------------------------------------------------- # Profile helpers # --------------------------------------------------------------------------- def load_profiles() -> list[dict]: """Scan collectors/profiles/*.yaml and return name+description dicts.""" profiles = [] if not PROFILES_DIR.is_dir(): return profiles try: import yaml except ImportError: # Fallback: parse name/description with basic string matching return _load_profiles_fallback() for yf in sorted(PROFILES_DIR.glob("*.yaml")): try: data = yaml.safe_load(yf.read_text()) if isinstance(data, dict): profiles.append({ "file": yf.name, "name": data.get("name", yf.stem), "description": data.get("description", ""), "base_url": data.get("base_url", ""), "expected_min_count": data.get("expected_min_count", ""), }) except Exception as exc: logger.warning("Failed to parse %s: %s", yf.name, exc) return profiles def _load_profiles_fallback() -> list[dict]: """Parse YAML profiles without the yaml library (best-effort).""" profiles = [] for yf in sorted(PROFILES_DIR.glob("*.yaml")): name = yf.stem description = "" base_url = "" for line in yf.read_text().splitlines(): line = line.strip() if line.startswith("name:"): name = line.split(":", 1)[1].strip().strip('"').strip("'") elif line.startswith("description:"): description = line.split(":", 1)[1].strip().strip('"').strip("'") elif line.startswith("base_url:"): base_url = line.split(":", 1)[1].strip().strip('"').strip("'") profiles.append({ "file": yf.name, "name": name, "description": description, "base_url": base_url, "expected_min_count": "", }) return profiles # --------------------------------------------------------------------------- # Click CLI # --------------------------------------------------------------------------- @click.group(invoke_without_command=True) @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--url", type=str, default=None, help="Ad-hoc source URL.") @click.option("--profile", "profile_path", type=click.Path(exists=True), default=None, help="Custom YAML profile path.") @click.option("--dry-run", is_flag=True, default=False, help="Estimate without running.") @click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed docs.") @click.option("--since", type=str, default=None, help="Only docs after DATE (YYYY-MM-DD).") @click.option("--date-range", nargs=2, type=str, default=None, help="Date window: START END (YYYY-MM-DD).") @click.option("--since-last-run", is_flag=True, default=False, help="Only new/changed since last run.") @click.option("--update", is_flag=True, default=False, help="Re-check cached URLs.") @click.option("--db", "db_path", type=click.Path(), default=None, help="Database path.") @click.option("--verbose", "-v", is_flag=True, default=False, help="Enable debug logging.") @click.pass_context def cli(ctx, source, url, profile_path, dry_run, retry_failed, since, date_range, since_last_run, update, db_path, verbose): """Mosaic -- Intelligence Extraction Platform.""" setup_logging(verbose) cfg = load_config() ctx.ensure_object(dict) ctx.obj["cfg"] = cfg ctx.obj["source"] = source ctx.obj["url"] = url ctx.obj["profile_path"] = profile_path ctx.obj["dry_run"] = dry_run ctx.obj["retry_failed"] = retry_failed ctx.obj["since"] = since ctx.obj["date_range"] = date_range ctx.obj["since_last_run"] = since_last_run ctx.obj["update"] = update ctx.obj["db_path"] = get_db_path(cfg, db_path) ctx.obj["verbose"] = verbose ctx.obj["output_dir"] = get_output_dir(cfg) if ctx.invoked_subcommand is None: interactive_menu(ctx) # --------------------------------------------------------------------------- # collect # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--url", type=str, default=None, help="Ad-hoc source URL.") @click.option("--name", type=str, default=None, help="Name for ad-hoc source.") @click.option("--profile", "profile_path", type=click.Path(exists=True), default=None, help="Custom YAML profile path.") @click.option("--limit", type=int, default=None, help="Maximum documents to fetch.") @click.option("--since", type=str, default=None, help="Only docs after DATE (YYYY-MM-DD).") @click.option("--date-range", nargs=2, type=str, default=None, help="Date window: START END.") @click.option("--update", is_flag=True, default=False, help="Re-check cached URLs.") @click.pass_context def collect(ctx, source, url, name, profile_path, limit, since, date_range, update): """Download documents from a source.""" source = source or ctx.obj.get("source") url = url or ctx.obj.get("url") if not source and not url and not profile_path: console.print("[yellow]No source specified.[/yellow] Use --source, --url, or --profile.") raise SystemExit(1) from collectors.profile_loader import load_profile, create_adhoc_profile from collectors.custom import CustomSource from utils.cache import DownloadCache from utils.http_limiter import HTTPRateLimiter # Load profile try: if profile_path: profile = load_profile(profile_path) elif url: pname = name or url.split('/')[2].replace('.', '_') profile = create_adhoc_profile(url, pname) else: profile = load_profile(source) except (FileNotFoundError, ValueError) as e: console.print(f"[red]Profile error: {e}[/red]") raise SystemExit(1) label = profile['name'] console.print(Panel( f"[bold cyan]Collect[/bold cyan]\n\n" f" Source : {label} — {profile.get('description', '')}\n" f" Limit : {limit or 'all'}\n" f" Since : {since or 'n/a'}\n" f" Range : {' to '.join(date_range) if date_range else 'n/a'}\n" f" Update : {'yes' if update else 'no'}", title="[bold]mosaic collect[/bold]", border_style="cyan", )) db = open_db(ctx.obj["db_path"]) try: cache = DownloadCache(os.path.join(ctx.obj["output_dir"], "raw")) cfg = ctx.obj["cfg"] http_cfg = { 'timeout': cfg.get('http', 'timeout', fallback='30'), 'max_retries': cfg.get('http', 'max_retries', fallback='3'), 'user_agent': cfg.get('http', 'user_agent', fallback='Mozilla/5.0 (compatible; research-crawler/1.0)'), } collector = CustomSource(profile, db, cache, config=http_cfg) since_date = since until_date = None if date_range: since_date, until_date = date_range stats = collector.collect( since=since_date, until=until_date, limit=limit, update=update, ) table = Table(title="Collection Results", box=box.ROUNDED) table.add_column("Metric", style="bold") table.add_column("Count", justify="right") table.add_row("Discovered", str(stats.get('discovered', 0))) table.add_row("Downloaded", str(stats.get('downloaded', 0))) table.add_row("Already cached", str(stats.get('cached', 0))) table.add_row("Failed", str(stats.get('failed', 0))) if update: table.add_row("Updated", str(stats.get('updated', 0))) console.print(table) finally: db.close() # --------------------------------------------------------------------------- # parse # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed docs.") @click.pass_context def parse(ctx, source, retry_failed): """Parse downloaded documents into structured text.""" source = source or ctx.obj.get("source") retry_failed = retry_failed or ctx.obj.get("retry_failed", False) from parsers.base import ParserManager from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn console.print(Panel( f"[bold green]Parse[/bold green]\n\n" f" Source : {source or 'all'}\n" f" Retry failed: {'yes' if retry_failed else 'no'}", title="[bold]mosaic parse[/bold]", border_style="green", )) db = open_db(ctx.obj["db_path"]) try: cfg = ctx.obj["cfg"] parser = ParserManager( memory_limit_mb=int(cfg.get('parser', 'subprocess_memory_limit_mb', fallback='512')), timeout_seconds=int(cfg.get('parser', 'subprocess_timeout_seconds', fallback='60')), min_chars_per_page=int(cfg.get('parser', 'min_chars_per_page', fallback='50')), max_pdf_pages=int(cfg.get('parser', 'max_pdf_pages', fallback='200')), ) # Get documents to parse if source: docs = db.get_documents_by_source(source, status='CACHED') else: docs = db.get_documents_by_source('', status='CACHED') # Get all cached docs across sources cursor = db.conn.execute( "SELECT * FROM documents WHERE status = 'CACHED' AND parse_status IN ('PENDING', ?)", ('PARSE_FAILED' if retry_failed else '__NONE__',) ) from storage.models import Document docs = [Document.from_row(row) for row in cursor.fetchall()] if not docs: console.print("[yellow]No documents to parse.[/yellow]") return parsed = 0 failed = 0 needs_ocr = 0 with Progress( SpinnerColumn(), TextColumn("{task.description}"), BarColumn(), TaskProgressColumn(), ) as progress: task = progress.add_task(f"Parsing {len(docs)} docs", total=len(docs)) for doc in docs: if doc.parse_status in ('PARSED', 'NEEDS_OCR') and not retry_failed: progress.advance(task) continue 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 if result.metadata: doc.metadata.update(result.metadata) ps = 'NEEDS_OCR' if result.needs_ocr else 'PARSED' doc.parse_status = ps db.insert_document(doc) db.update_document_status(doc.id, parse_status=ps) if result.needs_ocr: needs_ocr += 1 else: parsed += 1 else: db.update_document_status(doc.id, parse_status='PARSE_FAILED') failed += 1 progress.advance(task) table = Table(title="Parse Results", box=box.ROUNDED) table.add_column("Metric", style="bold") table.add_column("Count", justify="right") table.add_row("Parsed", str(parsed)) table.add_row("Needs OCR", str(needs_ocr)) table.add_row("Failed", str(failed)) console.print(table) finally: db.close() # --------------------------------------------------------------------------- # extract # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--domain", "-d", multiple=True, type=click.Choice(DOMAIN_CHOICES, case_sensitive=False), help="Extraction domain(s). Repeat for multiple.") @click.option("--dry-run", is_flag=True, default=False, help="Estimate only.") @click.option("--since", type=str, default=None, help="Only docs after DATE.") @click.option("--date-range", nargs=2, type=str, default=None, help="Date window.") @click.option("--since-last-run", is_flag=True, default=False, help="Only new/changed since last run.") @click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed.") @click.pass_context def extract(ctx, source, domain, dry_run, since, date_range, since_last_run, retry_failed): """Run AI-powered extraction on parsed documents.""" source = source or ctx.obj.get("source") dry_run = dry_run or ctx.obj.get("dry_run", False) domains = list(domain) if domain else DOMAIN_CHOICES if dry_run: _extract_dry_run(ctx.obj["db_path"], source, domains, since, date_range, since_last_run) return console.print(Panel( f"[bold magenta]Extract[/bold magenta]\n\n" f" Source : {source or 'all'}\n" f" Domains : {', '.join(domains)}\n" f" Since : {since or 'n/a'}\n" f" Range : {' to '.join(date_range) if date_range else 'n/a'}\n" f" Since last run : {'yes' if since_last_run else 'no'}\n" f" Retry failed : {'yes' if retry_failed else 'no'}", title="[bold]mosaic extract[/bold]", border_style="magenta", )) console.print( "\n[dim]Extraction not yet implemented — see extractors/[/dim]" ) def _extract_dry_run(db_path: str, source: Optional[str], domains: list[str], since: Optional[str], date_range: Optional[tuple], since_last_run: bool): """Estimate extraction scope without running anything.""" from utils.text_utils import estimate_tokens db = open_db(db_path) try: if date_range: docs = db.get_documents_by_date_range(source=source, since=date_range[0], until=date_range[1]) elif since: docs = db.get_documents_by_date_range(source=source, since=since) elif since_last_run: docs = db.get_unextracted_documents(source=source, since_last_run=True) else: docs = db.get_unextracted_documents(source=source) total_chars = sum(d.char_count for d in docs) total_tokens = sum(estimate_tokens(d.text) for d in docs if d.text) table = Table(title="Extraction Dry-Run Estimate", box=box.ROUNDED) table.add_column("Metric", style="bold") table.add_column("Value", justify="right") table.add_row("Documents", str(len(docs))) table.add_row("Total characters", f"{total_chars:,}") table.add_row("Estimated tokens", f"{total_tokens:,}") table.add_row("Domains", ", ".join(domains)) table.add_row("Est. API calls", str(len(docs) * len(domains))) console.print(table) finally: db.close() # --------------------------------------------------------------------------- # analyze # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--limit", type=int, default=None, help="Max documents per stage.") @click.pass_context def analyze(ctx, source, limit): """Full pipeline: collect -> parse -> extract -> report.""" source = source or ctx.obj.get("source") if not source: console.print("[yellow]No source specified.[/yellow] Use --source.") raise SystemExit(1) console.print(Panel( f"[bold yellow]Analyze — Full Pipeline[/bold yellow]\n\n" f" Source : {source}\n" f" Limit : {limit or 'all'}\n\n" f" Stages:\n" f" 1. Collect — download documents\n" f" 2. Parse — extract structured text\n" f" 3. Extract — AI-powered extraction\n" f" 4. Report — generate output", title="[bold]mosaic analyze[/bold]", border_style="yellow", )) console.print( "\n[dim]Full pipeline not yet implemented — run individual stages.[/dim]" ) # --------------------------------------------------------------------------- # query # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.option("--mode", type=click.Choice(QUERY_MODES, case_sensitive=False), default="default", help="Query mode.") @click.option("--actor", type=str, default=None, help="Filter to actor/group name.") @click.option("--topic", type=str, default=None, help="Filter to topic.") @click.pass_context def query(ctx, source, mode, actor, topic): """Interactive Q&A against the intelligence database.""" source = source or ctx.obj.get("source") db_path = ctx.obj["db_path"] db = open_db(db_path) try: _query_loop(db, source=source, mode=mode, actor=actor, topic=topic) finally: db.close() def _query_loop(db, source: Optional[str] = None, mode: str = "default", actor: Optional[str] = None, topic: Optional[str] = None): """Rich-based interactive query loop with FTS5 search.""" console.print(Panel( "[bold cyan]Mosaic Query Console[/bold cyan]\n\n" " Type a search query to find documents via full-text search.\n" " Special commands:\n" " [bold]help[/bold] — show this message\n" " [bold]exit[/bold] / [bold]quit[/bold] — leave the console\n" " [bold]mode [/bold] — switch query mode\n" " [bold]source [/bold] — filter by source", title="[bold]mosaic query[/bold]", border_style="cyan", )) if source: console.print(f" [dim]Source filter:[/dim] {source}") if mode != "default": console.print(f" [dim]Mode:[/dim] {mode}") if actor: console.print(f" [dim]Actor:[/dim] {actor}") if topic: console.print(f" [dim]Topic:[/dim] {topic}") console.print() current_source = source current_mode = mode while True: try: user_input = Prompt.ask("[bold cyan]mosaic[/bold cyan]") except (EOFError, KeyboardInterrupt): console.print("\n[dim]Exiting.[/dim]") break text = user_input.strip() if not text: continue # --- special commands --- lower = text.lower() if lower in ("exit", "quit", "q"): console.print("[dim]Exiting.[/dim]") break if lower == "help": console.print( " [bold][/bold] — full-text search\n" " [bold]mode [/bold] — switch mode (default/emulate/compare/brief/scenario)\n" " [bold]source [/bold] — set source filter (blank to clear)\n" " [bold]exit / quit[/bold] — leave" ) continue if lower.startswith("mode "): new_mode = text[5:].strip().lower() if new_mode in QUERY_MODES: current_mode = new_mode console.print(f" [dim]Mode set to:[/dim] {current_mode}") else: console.print(f" [red]Unknown mode.[/red] Choose from: {', '.join(QUERY_MODES)}") continue if lower.startswith("source "): new_source = text[7:].strip() current_source = new_source if new_source else None console.print(f" [dim]Source filter:[/dim] {current_source or '(none)'}") continue # --- FTS search --- results = db.search_documents(text, source=current_source, limit=20) if not results: console.print(" [dim]No results found.[/dim]") continue console.print(f"\n [bold]{len(results)} result(s)[/bold]\n") for i, doc in enumerate(results, 1): title = doc.title or doc.source_url or doc.id[:16] date_str = doc.date or "" src = doc.source or "" console.print( f" [bold cyan]{i}.[/bold cyan] {title}", markup=True, highlight=False, ) console.print( f" [dim]{src} {date_str} chars={doc.char_count}[/dim]" ) # Show a text snippet (first 300 chars) if doc.text: snippet = doc.text[:300].replace("\n", " ") console.print(f" {snippet}", markup=False, highlight=False) console.print() # --------------------------------------------------------------------------- # export # --------------------------------------------------------------------------- @cli.command(name="export") @click.option("--category", "-c", type=click.Choice(EXPORT_CATEGORIES, case_sensitive=False), default=None, help="Category to export (default: all).") @click.option("--format", "fmt", type=str, default="jsonl", help="Export format (default: jsonl).") @click.pass_context def export_cmd(ctx, category, fmt): """Export extracted intelligence to JSONL files.""" db_path = ctx.obj["db_path"] output_dir = ctx.obj["output_dir"] if fmt != "jsonl": console.print(f"[yellow]Format '{fmt}' not supported yet. Using jsonl.[/yellow]") db = open_db(db_path) try: store = open_jsonl_store(db, output_dir) with console.status("[bold green]Exporting..."): created = store.export(category=category) if created: table = Table(title="Exported Files", box=box.ROUNDED) table.add_column("File", style="bold") table.add_column("Path") for fpath in created: table.add_row(Path(fpath).name, fpath) console.print(table) else: console.print("[dim]Nothing to export — database is empty for the requested category.[/dim]") finally: db.close() # --------------------------------------------------------------------------- # ingest # --------------------------------------------------------------------------- MIME_TO_DOCTYPE = { 'application/pdf': 'pdf', 'text/html': 'html', 'text/plain': 'text', 'message/rfc822': 'email', } def _detect_doc_type(file_path: str) -> str: """Detect document type from file content.""" from utils.sanitize import validate_content_type mime = validate_content_type(file_path) return MIME_TO_DOCTYPE.get(mime, 'text') def _hash_file(file_path: str) -> str: """Compute SHA256 of a file.""" sha = hashlib.sha256() with open(file_path, 'rb') as f: for chunk in iter(lambda: f.read(8192), b''): sha.update(chunk) return sha.hexdigest() def _ingest_raw_file(db, cache, file_path: str, source_name: str) -> bool: """Ingest a single raw document file into the DB and cache. Returns True if the document was new, False if already exists. """ from storage.models import Document fpath = Path(file_path) content_hash = _hash_file(str(fpath)) existing = db.get_document(content_hash) if existing: return False doc_type = _detect_doc_type(str(fpath)) ext = fpath.suffix.lstrip('.') # Store in cache cached_path = cache.get_shard_path(content_hash, source_name, ext=ext) cached_path.parent.mkdir(parents=True, exist_ok=True) import shutil shutil.copy2(str(fpath), str(cached_path)) doc = Document( id=content_hash, source=source_name, source_url=str(fpath.resolve()), doc_type=doc_type, title=fpath.stem, raw_path=str(cached_path), content_hash=content_hash, fetch_date=datetime.now().isoformat(), char_count=fpath.stat().st_size, status='CACHED', parse_status='PENDING', ) db.insert_document(doc) return True @cli.command() @click.option("--file", "-f", "files", multiple=True, type=click.Path(exists=True), help="File(s) to ingest. Repeat for multiple files.") @click.option("--dir", "-d", "directory", type=click.Path(exists=True, file_okay=False), help="Directory of files to ingest.") @click.option("--source", "-s", type=str, required=True, help="Source name to tag ingested docs.") @click.option("--category", "-c", type=click.Choice(EXPORT_CATEGORIES, case_sensitive=False), default=None, help="JSONL category (for importing pre-extracted data).") @click.option("--watch", "-w", is_flag=True, default=False, help="Watch directory for new files (real-time mode).") @click.option("--recursive", "-r", is_flag=True, default=False, help="Recurse into subdirectories.") @click.option("--parse", "auto_parse", is_flag=True, default=False, help="Automatically parse ingested documents.") @click.pass_context def ingest(ctx, files, directory, source, category, watch, recursive, auto_parse): """Ingest local files or data dumps into the Mosaic pipeline. Import raw documents (PDF, HTML, text) or pre-extracted JSONL data without HTTP collection. Supports single files, directories, and real-time directory watching. \b Examples: mosaic ingest -f dump.pdf -s wikileaks_dump mosaic ingest -d /path/to/cables/ -s cablegate -r --parse mosaic ingest -f extracted.jsonl -s vault7 -c tools mosaic ingest -d /incoming/ -s live_feed -w """ db_path = ctx.obj["db_path"] output_dir = ctx.obj["output_dir"] db = open_db(db_path) try: # JSONL import mode if category: store = open_jsonl_store(db, output_dir) total = 0 targets = list(files) if directory: p = Path(directory) targets.extend(str(f) for f in p.glob("*.jsonl")) if not targets: console.print("[yellow]No JSONL files to import.[/yellow]") return for fpath in targets: with console.status(f"[bold green]Importing {Path(fpath).name}..."): count = store.import_jsonl(fpath, category) console.print(f" [green]+{count}[/green] {category} from {Path(fpath).name}") total += count console.print(f"\n[bold green]Imported {total} {category} records total.[/bold green]") return # Raw document ingest mode cache = open_cache(output_dir) file_list = list(files) if directory: p = Path(directory) pattern = "**/*" if recursive else "*" file_list.extend( str(f) for f in p.glob(pattern) if f.is_file() and not f.name.startswith('.') ) if watch and directory: _watch_directory(db, cache, directory, source, recursive, auto_parse, ctx) return if not file_list: console.print("[yellow]No files to ingest.[/yellow] Use --file, --dir, or --watch.") return console.print(Panel( f"[bold green]Ingest[/bold green]\n\n" f" Source : {source}\n" f" Files : {len(file_list)}\n" f" Parse : {'yes' if auto_parse else 'no'}", title="[bold]mosaic ingest[/bold]", border_style="green", )) new_count = 0 skip_count = 0 with Progress( SpinnerColumn(), TextColumn("{task.description}"), BarColumn(), TextColumn("{task.completed}/{task.total}"), ) as progress: task = progress.add_task("Ingesting files", total=len(file_list)) for fpath in file_list: is_new = _ingest_raw_file(db, cache, fpath, source) if is_new: new_count += 1 else: skip_count += 1 progress.advance(task) table = Table(title="Ingest Results", box=box.ROUNDED) table.add_column("Metric", style="bold") table.add_column("Count", justify="right") table.add_row("[green]New documents[/green]", str(new_count)) table.add_row("[dim]Already cached[/dim]", str(skip_count)) table.add_row("[bold]Total processed[/bold]", str(new_count + skip_count)) console.print(table) # Auto-parse if requested if auto_parse and new_count > 0: console.print("\n[bold]Running parser on new documents...[/bold]") ctx.invoke(parse, source=source) finally: db.close() def _watch_directory(db, cache, directory: str, source: str, recursive: bool, auto_parse: bool, ctx): """Watch a directory for new files and ingest them in real-time.""" import time as _time watched = Path(directory) seen = set() # Seed with existing files pattern = "**/*" if recursive else "*" for f in watched.glob(pattern): if f.is_file() and not f.name.startswith('.'): seen.add(str(f.resolve())) console.print(Panel( f"[bold cyan]Watch Mode[/bold cyan]\n\n" f" Directory : {directory}\n" f" Source : {source}\n" f" Recursive : {'yes' if recursive else 'no'}\n" f" Auto-parse: {'yes' if auto_parse else 'no'}\n\n" f" [dim]Press Ctrl+C to stop watching.[/dim]", title="[bold]mosaic ingest --watch[/bold]", border_style="cyan", )) console.print(f" [dim]Watching {len(seen)} existing files (skipped)...[/dim]") try: while True: new_files = [] for f in watched.glob(pattern): resolved = str(f.resolve()) if f.is_file() and not f.name.startswith('.') and resolved not in seen: new_files.append(f) seen.add(resolved) for fpath in new_files: is_new = _ingest_raw_file(db, cache, str(fpath), source) if is_new: console.print(f" [green]+[/green] {fpath.name}") if auto_parse: from parsers.base import ParserManager from storage.models import Document cfg = ctx.obj["cfg"] parser = ParserManager( memory_limit_mb=int(cfg.get('parser', 'subprocess_memory_limit_mb', fallback='512')), timeout_seconds=int(cfg.get('parser', 'subprocess_timeout_seconds', fallback='60')), ) doc = db.get_document(_hash_file(str(fpath))) if doc and doc.parse_status == 'PENDING': 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 ps = 'NEEDS_OCR' if result.needs_ocr else 'PARSED' db.insert_document(doc) db.update_document_status(doc.id, parse_status=ps) console.print(f" [dim]parsed → {ps}[/dim]") else: db.update_document_status(doc.id, parse_status='PARSE_FAILED') console.print(f" [red]parse failed[/red]") _time.sleep(2) except KeyboardInterrupt: console.print("\n[dim]Watch stopped.[/dim]") # --------------------------------------------------------------------------- # status # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.pass_context def status(ctx, source): """Show progress dashboard with collection, parse, and extraction stats.""" source = source or ctx.obj.get("source") db_path = ctx.obj["db_path"] db = open_db(db_path) try: _show_status(db, source) finally: db.close() def _show_status(db, source: Optional[str] = None): """Render status tables.""" counts = db.get_status_counts(source=source) ext_counts = db.get_extraction_counts(source=source) tokens = db.get_total_tokens(source=source) title_suffix = f" [{source}]" if source else "" # --- Collection status --- coll_table = Table(title=f"Collection Status{title_suffix}", box=box.ROUNDED) coll_table.add_column("Status", style="bold") coll_table.add_column("Count", justify="right") coll_statuses = ["PENDING", "DOWNLOADING", "CACHED", "DOWNLOAD_FAILED"] for s in coll_statuses: cnt = counts["collection"].get(s, 0) style = "green" if s == "CACHED" else ("red" if "FAIL" in s else "") coll_table.add_row(s, f"[{style}]{cnt}[/{style}]" if style else str(cnt)) coll_table.add_row("[bold]Total[/bold]", f"[bold]{counts['total']}[/bold]") # --- Parse status --- parse_table = Table(title=f"Parse Status{title_suffix}", box=box.ROUNDED) parse_table.add_column("Status", style="bold") parse_table.add_column("Count", justify="right") parse_statuses = ["PENDING", "IN_PROGRESS", "PARSED", "PARSE_FAILED", "NEEDS_OCR"] for s in parse_statuses: cnt = counts["parsing"].get(s, 0) style = "green" if s == "PARSED" else ("red" if "FAIL" in s else ("yellow" if s == "NEEDS_OCR" else "")) parse_table.add_row(s, f"[{style}]{cnt}[/{style}]" if style else str(cnt)) # --- Extraction status per extractor --- ext_table = Table(title=f"Extraction Status{title_suffix}", box=box.ROUNDED) ext_table.add_column("Extractor", style="bold") ext_table.add_column("PENDING", justify="right") ext_table.add_column("IN_PROGRESS", justify="right") ext_table.add_column("DONE", justify="right", style="green") ext_table.add_column("FAILED", justify="right", style="red") if ext_counts: for ext_name, statuses in sorted(ext_counts.items()): ext_table.add_row( ext_name, str(statuses.get("PENDING", 0)), str(statuses.get("IN_PROGRESS", 0)), str(statuses.get("DONE", 0)), str(statuses.get("FAILED", 0)), ) else: ext_table.add_row("[dim]no extractions yet[/dim]", "", "", "", "") # --- Token usage --- token_table = Table(title=f"Token Usage{title_suffix}", box=box.ROUNDED) token_table.add_column("Metric", style="bold") token_table.add_column("Value", justify="right") token_table.add_row("Input tokens", f"{tokens['input_tokens']:,}") token_table.add_row("Output tokens", f"{tokens['output_tokens']:,}") token_table.add_row("[bold]Total tokens[/bold]", f"[bold]{tokens['total_tokens']:,}[/bold]") console.print() console.print(coll_table) console.print() console.print(parse_table) console.print() console.print(ext_table) console.print() console.print(token_table) console.print() # --------------------------------------------------------------------------- # verify # --------------------------------------------------------------------------- @cli.command() @click.option("--source", type=str, default=None, help="Source profile name.") @click.pass_context def verify(ctx, source): """Re-hash cached files and report integrity status.""" source = source or ctx.obj.get("source") db_path = ctx.obj["db_path"] output_dir = ctx.obj["output_dir"] db = open_db(db_path) try: _run_verify(db, output_dir, source) finally: db.close() def _run_verify(db, output_dir: str, source: Optional[str] = None): """Verify cached files against stored hashes.""" cache = open_cache(output_dir) doc_hashes = db.get_all_document_hashes() if not doc_hashes: console.print("[dim]No cached documents to verify.[/dim]") return verified = 0 mismatched = 0 missing = 0 bad_files: list[str] = [] with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TextColumn("{task.completed}/{task.total}"), console=console, ) as progress: task = progress.add_task("Verifying files...", total=len(doc_hashes)) for doc_id, raw_path, expected_hash in doc_hashes: if source: # Filter by source if provided — look up in DB doc = db.get_document(doc_id) if doc and doc.source != source: progress.advance(task) continue if not raw_path or not Path(raw_path).exists(): missing += 1 progress.advance(task) continue actual_hash = cache.compute_hash(raw_path) if actual_hash == expected_hash: verified += 1 else: mismatched += 1 bad_files.append(raw_path) progress.advance(task) # Results table table = Table(title="Integrity Verification Results", box=box.ROUNDED) table.add_column("Status", style="bold") table.add_column("Count", justify="right") table.add_row("[green]Verified[/green]", str(verified)) table.add_row("[red]Mismatched[/red]", str(mismatched)) table.add_row("[yellow]Missing[/yellow]", str(missing)) table.add_row("[bold]Total checked[/bold]", str(verified + mismatched + missing)) console.print() console.print(table) if bad_files: console.print("\n[red bold]Mismatched files:[/red bold]") for bf in bad_files: console.print(f" [red]- {bf}[/red]", markup=False) elif mismatched == 0 and missing == 0: console.print("\n[green]All cached files passed integrity checks.[/green]") # --------------------------------------------------------------------------- # manual # --------------------------------------------------------------------------- @cli.command() @click.option("--output", "-o", type=click.Path(), default="output/FIELD_MANUAL.md", help="Output file path.") @click.option("--title", type=str, default="Operator's Field Manual", help="Manual title.") @click.pass_context def manual(ctx, output, title): """Generate Markdown field manual from extracted intelligence.""" from analysis.field_manual import generate_field_manual db = open_db(ctx.obj["db_path"]) try: path = generate_field_manual(db, output_path=output, title=title) console.print(f"\n[bold green]Field manual generated:[/bold green] {path}") # Show stats import os size = os.path.getsize(path) with open(path) as f: lines = sum(1 for _ in f) console.print(f" Size: {size:,} bytes ({lines:,} lines)") finally: db.close() # --------------------------------------------------------------------------- # sources # --------------------------------------------------------------------------- @cli.command() @click.pass_context def sources(ctx): """List available source profiles.""" profiles = load_profiles() table = Table(title="Source Profiles", box=box.ROUNDED) table.add_column("Name", style="bold cyan") table.add_column("Description") table.add_column("Base URL", style="dim") table.add_column("Expected Docs", justify="right") table.add_column("File", style="dim") if profiles: for p in profiles: table.add_row( p["name"], p["description"], p["base_url"], str(p["expected_min_count"]) if p["expected_min_count"] else "", p["file"], ) else: table.add_row("[dim]no profiles found[/dim]", "", "", "", "") console.print() console.print(table) console.print(f"\n [dim]Profile directory: {PROFILES_DIR}[/dim]") console.print(" [dim]Use --profile to load a custom YAML file from any path.[/dim]\n") # --------------------------------------------------------------------------- # Interactive menu # --------------------------------------------------------------------------- MENU_ITEMS = [ ("1", "Collect", "Download documents", "collect"), ("2", "Parse", "Extract structured text", "parse"), ("3", "Extract", "AI-powered extraction", "extract"), ("4", "Analyze", "Full pipeline", "analyze"), ("5", "Query", "Interactive Q&A", "query"), ("6", "Export", "Export to JSONL", "export"), ("7", "Ingest", "Import local files/dumps", "ingest"), ("8", "Status", "Progress dashboard", "status"), ("9", "Verify", "Integrity check", "verify"), ("s", "Sources", "List profiles", "sources"), ("0", "Exit", "", None), ] def interactive_menu(ctx: click.Context): """Display the main interactive menu and dispatch commands.""" while True: lines = [] for key, label, desc, _ in MENU_ITEMS: if key == "0": lines.append(f" [bold red][{key}][/bold red] {label}") else: lines.append(f" [bold cyan][{key}][/bold cyan] {label:<12} [dim]{desc}[/dim]") panel_text = "\n".join(lines) console.print() console.print(Panel( panel_text, title="[bold white]MOSAIC[/bold white] [dim]Intelligence Extraction Platform[/dim]", border_style="bright_blue", padding=(1, 2), )) try: choice = Prompt.ask("\n[bold bright_blue]Select[/bold bright_blue]", choices=[m[0] for m in MENU_ITEMS], show_choices=False) except (EOFError, KeyboardInterrupt): console.print("\n[dim]Exiting.[/dim]") break if choice == "0": console.print("[dim]Exiting.[/dim]") break # Find the matching menu item for key, label, desc, cmd_name in MENU_ITEMS: if key == choice and cmd_name: _dispatch_interactive(ctx, cmd_name) break def _dispatch_interactive(ctx: click.Context, cmd_name: str): """Invoke a subcommand from the interactive menu.""" cmd = cli.get_command(ctx, cmd_name) if cmd is None: console.print(f"[red]Unknown command: {cmd_name}[/red]") return # Some commands benefit from prompting for a source first if cmd_name in ("collect", "parse", "extract", "analyze", "status", "verify"): source = _prompt_source_optional() if source: ctx.obj["source"] = source try: sub_ctx = click.Context(cmd, parent=ctx, info_name=cmd_name) sub_ctx.ensure_object(dict) sub_ctx.obj = ctx.obj cmd.invoke(sub_ctx) except SystemExit: pass except Exception as exc: console.print(f"[red]Error: {exc}[/red]") logger.exception("Command '%s' failed", cmd_name) def _prompt_source_optional() -> Optional[str]: """Optionally prompt for a source profile name.""" profiles = load_profiles() if profiles: names = [p["name"] for p in profiles] console.print(f"\n [dim]Available sources: {', '.join(names)}[/dim]") source = Prompt.ask( " [dim]Source (blank for all)[/dim]", default="", show_default=False, ) return source.strip() or None # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- def main(): """Entry point for mosaic CLI.""" # Ensure we can import project modules regardless of cwd if str(BASE_DIR) not in sys.path: sys.path.insert(0, str(BASE_DIR)) cli(standalone_mode=True) if __name__ == "__main__": main()