Add main CLI with Click subcommands and interactive menu
This commit is contained in:
n0mad1k
@@ -0,0 +1,857 @@
|
||||
#!/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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MosaicCLI(click.Group):
|
||||
"""Custom group that shows the interactive menu when invoked bare."""
|
||||
|
||||
def invoke(self, ctx):
|
||||
if not ctx.protected_args and not ctx.invoked_subcommand:
|
||||
interactive_menu(ctx)
|
||||
else:
|
||||
super().invoke(ctx)
|
||||
|
||||
|
||||
@click.group(cls=MosaicCLI, 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)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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)
|
||||
|
||||
label = source or url or profile_path
|
||||
console.print(Panel(
|
||||
f"[bold cyan]Collect[/bold cyan]\n\n"
|
||||
f" Source : {label}\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",
|
||||
))
|
||||
console.print(
|
||||
"\n[dim]Collection not yet implemented — see collectors/[/dim]"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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)
|
||||
|
||||
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",
|
||||
))
|
||||
console.print(
|
||||
"\n[dim]Parsing not yet implemented — see parsers/[/dim]"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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 <name>[/bold] — switch query mode\n"
|
||||
" [bold]source <name>[/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]<query>[/bold] — full-text search\n"
|
||||
" [bold]mode <name>[/bold] — switch mode (default/emulate/compare/brief/scenario)\n"
|
||||
" [bold]source <name>[/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()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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]")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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", "Status", "Progress dashboard", "status"),
|
||||
("8", "Verify", "Integrity check", "verify"),
|
||||
("9", "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()
|
||||
Reference in New Issue
Block a user