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())))292293294if __name__ == "__main__":295 app()296