HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""`aia` — AI Atlas operations CLI."""2from __future__ import annotations34import asyncio5import json6import logging7from pathlib import Path8from typing import Annotated910import typer11from rich.console import Console12from rich.table import Table1314from aiatlas.config import settings15from aiatlas.logging import setup_logging1617app = typer.Typer(name="aia", help="AI Atlas — the global intelligence layer for AI.", no_args_is_help=True, add_completion=False)18console = Console(stderr=True)19out = Console()202122@app.callback()23def _main(verbose: Annotated[bool, typer.Option("--verbose", "-v")] = False) -> None:24 setup_logging(level=logging.DEBUG if verbose else logging.INFO, service="aia-cli")25 settings.ensure_dirs()262728def _run(coro): # type: ignore[no-untyped-def]29 from aiatlas.db import dispose30 from aiatlas.services import cache3132 async def wrapper(): # type: ignore[no-untyped-def]33 try:34 return await coro35 finally:36 await dispose()37 await cache.close()3839 return asyncio.run(wrapper())404142@app.command()43def migrate(revision: str = "head") -> None:44 """Apply database migrations (forward-only)."""45 from alembic import command46 from alembic.config import Config4748 root = Path(__file__).resolve().parents[2]49 cfg = Config(str(root / "alembic.ini"))50 cfg.set_main_option("script_location", str(root / "migrations"))51 command.upgrade(cfg, revision)52 out.print("[green]migrations applied[/]")535455@app.command()56def seed() -> None:57 """Seed sources, connectors and curated registries (idempotent)."""58 from aiatlas.db import transaction59 from aiatlas.registry.seed import seed as _seed6061 async def go(): # type: ignore[no-untyped-def]62 async with transaction() as conn:63 return await _seed(conn)6465 res = _run(go())66 out.print(f"[green]seeded[/] {json.dumps(res)}")676869@app.command()70def connectors() -> None:71 """List registered connectors (code) and their database state."""72 from aiatlas.connectors import registry73 from aiatlas.db import fetch_all, transaction7475 async def go(): # type: ignore[no-untyped-def]76 async with transaction() as conn:77 return await fetch_all(conn, "select * from connectors order by priority, name")7879 rows = {r["name"]: r for r in _run(go())}80 t = Table(title="connectors")81 for col in ("name", "tier", "enabled", "health", "interval", "last success", "last change", "failures", "parser"):82 t.add_column(col)83 for name, cls in registry().items():84 r = rows.get(name)85 t.add_row(name, str(cls.tier), "✓" if (r and r["enabled"]) else "✗", r["health"] if r else "—", f"{(r['interval_seconds'] if r else cls.interval_seconds) // 60} min",86 r["last_success_at"].strftime("%m-%d %H:%M") if r and r["last_success_at"] else "—",87 r["last_change_at"].strftime("%m-%d %H:%M") if r and r["last_change_at"] else "—",88 str(r["consecutive_failures"]) if r else "—", cls.parser_version)89 out.print(t)909192@app.command()93def run(name: str, force: bool = typer.Option(False, "--force", help="ignore enabled/circuit/conditional headers"),94 file: list[str] = typer.Option(None, "--file", help="key=path or url=path local override (fixtures / seed snapshots)"),95 max_targets: int = typer.Option(0, "--max-targets"), url: list[str] = typer.Option(None, "--url", help="only these URLs")) -> None:96 """Run one connector now."""97 from aiatlas.connectors import get9899 overrides = dict(f.split("=", 1) for f in (file or []))100 ctx = _run(get(name).run(force=force, file_overrides=overrides, max_targets=max_targets or None, only_urls=url or None))101 out.print(json.dumps({k: v for k, v in ctx.stats.__dict__.items()}, default=str, indent=1))102103104@app.command()105def reprocess(name: str, url: list[str] = typer.Option(None, "--url")) -> None:106 """Re-extract from stored snapshots (no network) — after a parser improvement."""107 from aiatlas.connectors import get108109 ctx = _run(get(name).run(reprocess=True, force=True, only_urls=url or None))110 out.print(json.dumps({k: v for k, v in ctx.stats.__dict__.items()}, default=str, indent=1))111112113@app.command()114def crawl(priority: int = typer.Option(9, "--priority", help="run connectors with priority <= this"), force: bool = False,115 only: list[str] = typer.Option(None, "--only")) -> None:116 """Run every enabled connector once, in priority order (initial corpus build)."""117 from aiatlas.connectors import get, registry118 from aiatlas.db import fetch_all, transaction119120 async def go(): # type: ignore[no-untyped-def]121 async with transaction() as conn:122 rows = await fetch_all(conn, "select name, enabled, priority from connectors order by priority, name")123 results = {}124 for r in rows:125 if r["name"] not in registry() or (not r["enabled"] and not force) or r["priority"] > priority or (only and r["name"] not in only):126 continue127 console.print(f"[cyan]▶ {r['name']}[/]")128 try:129 ctx = await get(r["name"]).run(force=force)130 results[r["name"]] = {k: v for k, v in ctx.stats.__dict__.items() if k != "meta"}131 except Exception as exc: # noqa: BLE001132 results[r["name"]] = {"error": str(exc)[:300]}133 return results134135 out.print(json.dumps(_run(go()), default=str, indent=1))136137138@app.command()139def status() -> None:140 """Connector health table and queue depth."""141 from aiatlas.db import fetch_all, transaction142 from aiatlas.services.jobs import queue_depth143144 async def go(): # type: ignore[no-untyped-def]145 async with transaction() as conn:146 rows = await fetch_all(conn, """select c.name, c.enabled, c.health, c.last_success_at, c.next_run_at, c.consecutive_failures,147 (select count(*) from documents d where d.connector_name = c.name) as docs,148 (select count(*) from snapshots s join documents d on d.id = s.document_id where d.connector_name = c.name) as snaps,149 (select status from connector_runs r where r.connector_name = c.name order by started_at desc limit 1) as last_status150 from connectors c order by c.priority, c.name""")151 return rows, await queue_depth(conn)152153 rows, depth = _run(go())154 t = Table(title="AI Atlas connectors")155 for col in ("connector", "on", "health", "last run", "last success", "next run", "docs", "snapshots", "fails"):156 t.add_column(col)157 for r in rows:158 t.add_row(r["name"], "✓" if r["enabled"] else "✗", r["health"], r["last_status"] or "—",159 r["last_success_at"].strftime("%m-%d %H:%M") if r["last_success_at"] else "—",160 r["next_run_at"].strftime("%m-%d %H:%M") if r["next_run_at"] else "—", str(r["docs"]), str(r["snaps"]), str(r["consecutive_failures"]))161 out.print(t)162 out.print(f"queue: {json.dumps(depth)}")163164165@app.command()166def stats() -> None:167 """Live counters (from the database) and archive size."""168 from aiatlas.services.stats import compute_stats169170 out.print(json.dumps(_run(compute_stats()), default=str, indent=1))171172173@app.command()174def quality(limit: int = 20000) -> None:175 """Recompute entity quality scores."""176 from aiatlas.services.quality import recompute177178 out.print(json.dumps(_run(recompute(limit=limit))))179180181@app.command()182def schedule(no_worker: bool = typer.Option(False, "--no-worker")) -> None:183 """Long-running scheduler: due connectors, job worker, hourly stats/quality, nightly backup."""184 from aiatlas.services.scheduler import main185186 setup_logging(service="aia-scheduler")187 asyncio.run(main(with_worker=not no_worker))188189190@app.command()191def worker(concurrency: int = typer.Option(0, "--concurrency"), kind: list[str] = typer.Option(None, "--kind")) -> None:192 """Job worker only (LLM extraction, embeddings, reprocessing). Run on any node with database access."""193 from aiatlas.services.jobs import run_worker194195 setup_logging(service="aia-worker")196 asyncio.run(run_worker(concurrency=concurrency or None, kinds=kind or None))197198199@app.command()200def api(host: str = "", port: int = 0, reload: bool = False, workers: int = 1) -> None:201 """Serve the FastAPI application."""202 import uvicorn203204 uvicorn.run("aiatlas.api.main:app", host=host or settings.api_host, port=port or settings.api_port, reload=reload, workers=workers if not reload else 1,205 log_level="info", access_log=False, proxy_headers=True)206207208@app.command()209def backup() -> None:210 """pg_dump into AIA_DATA_DIR/backups (keeps 30)."""211 from aiatlas.services.backup import backup_database212213 out.print(str(backup_database()))214215216@app.command()217def llm(text: str = typer.Argument("", help="optional text to classify"), health: bool = typer.Option(True)) -> None:218 """Check the local LLM factory gateway (and optionally classify a snippet)."""219 from aiatlas.services.llm import gateway220221 async def go(): # type: ignore[no-untyped-def]222 h = await gateway.health()223 if text:224 h["classification"] = await gateway.classify(text=text, labels=["model_release", "pricing", "research_paper", "company_news", "other"])225 return h226227 out.print(json.dumps(_run(go()), indent=1))228229230@app.command()231def embed(limit: int = 200) -> None:232 """Embed entities missing vectors (requires the LLM gateway)."""233 from aiatlas.services.embeddings import embed_entities, pending_entity_ids234235 async def go(): # type: ignore[no-untyped-def]236 ids = await pending_entity_ids(limit)237 return await embed_entities(ids) if ids else {"embedded": 0}238239 out.print(json.dumps(_run(go())))240241242@app.command()243def search(q: str, limit: int = 10) -> None:244 """Search entities (natural-language filters are compiled deterministically)."""245 from aiatlas.db import transaction246 from aiatlas.services.search import compile_query, search_entities247248 async def go(): # type: ignore[no-untyped-def]249 query = compile_query(q)250 async with transaction() as conn:251 return query.as_dict(), await search_entities(conn, query, limit=limit)252253 query, rows = _run(go())254 out.print(json.dumps(query))255 for r in rows:256 out.print(f" {r['entity_type']:<10} {r['canonical_name']:<50} {r.get('organization_name') or ''} /{r['slug']}")257258259@app.command()260def review(status: str = "pending", limit: int = 30) -> None:261 """Show the human review queue."""262 from aiatlas.db import fetch_all, transaction263264 async def go(): # type: ignore[no-untyped-def]265 async with transaction() as conn:266 return await fetch_all(conn, "select id, kind, reason, created_at from review_queue where status = :s order by created_at desc limit :n", s=status, n=limit)267268 for r in _run(go()):269 out.print(f"{r['created_at']:%m-%d %H:%M} {r['kind']:<16} {r['reason']}")270271272@app.command()273def enqueue_llm(limit: int = 500, task: str = "auto") -> None:274 """Queue LLM extraction for snapshots that are still `llm_pending`/`stored` on documents flagged needs_llm."""275 from aiatlas.db import fetch_all, transaction276 from aiatlas.services.jobs import enqueue277278 async def go(): # type: ignore[no-untyped-def]279 async with transaction() as conn:280 rows = await fetch_all(conn, """select s.id, d.entity_id, d.connector_name from snapshots s join documents d on d.id = s.document_id281 where s.processing_status in ('llm_pending','stored') and d.needs_llm and s.changed282 and s.id = (select id from snapshots s2 where s2.document_id = s.document_id and s2.changed order by observed_at desc limit 1)283 order by s.observed_at desc limit :n""", n=limit)284 n = 0285 for r in rows:286 if await enqueue(conn, "llm_extract", {"snapshot_id": r["id"], "task": task, "entity_id": r["entity_id"], "connector": r["connector_name"]},287 priority=5, dedupe_key=f"llm_extract:{r['id']}"):288 n += 1289 return {"queued": n, "candidates": len(rows)}290291 out.print(json.dumps(_run(go())))292293294@app.command()295def canonicalize(apply: bool = typer.Option(False, "--apply", help="write changes (default: dry-run report)"),296 step: list[str] = typer.Option(None, "--step", help="run only these steps (repeatable); default = all, in canonical order"),297 as_json: bool = typer.Option(False, "--json", help="machine-readable report")) -> None:298 """Canonicalize the corpus: duplicates → effort variants → artifacts → families → licences → taxonomy → results → events → anomalies.299 Dry-run by default; idempotent under --apply (a second pass reports 0 changes). Never deletes rows, never touches snapshots."""300 from aiatlas.services.canonical import canonicalize as _canon301302 report = _run(_canon(apply=apply, steps=step or None))303 out.print(json.dumps(report.as_dict(), indent=1, default=str) if as_json else report.render())304305306@app.command()307def anomalies(severity: str = typer.Option("", "--severity", help="critical|warning|info"), status: str = typer.Option("open", "--status"),308 limit: int = typer.Option(100, "--limit"), refresh: bool = typer.Option(False, "--refresh", help="re-run the checks before listing")) -> None:309 """List data anomalies (flags with evidence — nothing is deleted or silently fixed)."""310 from aiatlas.db import transaction311 from aiatlas.services.anomalies import list_anomalies, run_checks312313 async def go(): # type: ignore[no-untyped-def]314 async with transaction() as conn:315 summary = await run_checks(conn) if refresh else None316 return summary, await list_anomalies(conn, severity=severity or None, status=status, limit=limit)317318 summary, rows = _run(go())319 if summary:320 console.print(f"[cyan]checks:[/] {json.dumps(summary)}")321 t = Table(title=f"anomalies ({status}{', ' + severity if severity else ''})")322 for col in ("severity", "check", "entity", "message", "last seen"):323 t.add_column(col)324 for r in rows:325 t.add_row(r["severity"], r["check_name"], r["slug"] or (r["entity_id"] or "—"), r["message"][:110], r["last_seen_at"].strftime("%m-%d %H:%M"))326 out.print(t)327328329quarantine_app = typer.Typer(help="Runs held by the anomaly detector (nothing written until released).", no_args_is_help=True)330app.add_typer(quarantine_app, name="quarantine")331332333@quarantine_app.command("list")334def quarantine_list(status: str = typer.Option("pending", "--status"), limit: int = 50) -> None:335 """Show quarantined runs."""336 from aiatlas.db import transaction337 from aiatlas.services.canonical import list_quarantine338339 async def go(): # type: ignore[no-untyped-def]340 async with transaction() as conn:341 return await list_quarantine(conn, status=status, limit=limit)342343 t = Table(title=f"quarantined runs ({status or 'all'})")344 for col in ("id", "connector", "created", "status", "docs", "reason"):345 t.add_column(col)346 for r in _run(go()):347 t.add_row(r["id"], r["connector_name"], r["created_at"].strftime("%m-%d %H:%M"), r["status"], str(r["documents"]), r["reason"][:90])348 out.print(t)349350351@quarantine_app.command("release")352def quarantine_release(quarantine_id: str) -> None:353 """Write the held facts as the connector would have."""354 from aiatlas.db import transaction355 from aiatlas.services.canonical import release_quarantine356357 async def go(): # type: ignore[no-untyped-def]358 async with transaction() as conn:359 return await release_quarantine(conn, quarantine_id, actor="cli")360361 out.print(json.dumps(_run(go()), default=str))362363364@quarantine_app.command("discard")365def quarantine_discard(quarantine_id: str, note: str = typer.Option("", "--note")) -> None:366 """Drop the held facts (the raw snapshots stay archived)."""367 from aiatlas.db import transaction368 from aiatlas.services.canonical import discard_quarantine369370 async def go(): # type: ignore[no-untyped-def]371 async with transaction() as conn:372 return await discard_quarantine(conn, quarantine_id, actor="cli", note=note or None)373374 out.print(json.dumps(_run(go()), default=str))375376377if __name__ == "__main__":378 app()379