spb/company-atlas

1"""`catlas` intelligence commands: process-changes, enrich, metrics, daily, signals, trends, alerts-eval, llm-test, reprocess-events, events, digest, retention."""2from __future__ import annotations34import asyncio5import json6import time7from datetime import UTC, date, datetime, timedelta8from typing import Annotated910import typer11from rich.table import Table1213from companyatlas.cli import console, out, run_async141516def register(app: typer.Typer) -> None:17    @app.command("process-changes")18    def process_changes(limit: int = 200, loop: Annotated[bool, typer.Option("--loop", help="keep polling every 20 s")] = False) -> None:19        """Turn pending changes into deterministic events (clusters, sources, review queue, LLM jobs)."""20        from companyatlas.services.events import process_pending_changes2122        async def go() -> None:23            while True:24                stats = await process_pending_changes(limit=limit)25                out.print(stats)26                if not loop:27                    return28                await asyncio.sleep(20)2930        run_async(go())3132    @app.command()33    def enrich(limit: int = 10, once: Annotated[bool, typer.Option("--once", help="one batch only (default: drain until idle)")] = False) -> None:34        """Run the LLM enrichment worker over pending llm_jobs (sticky by kind, budgeted)."""35        from companyatlas.services.llm.enrich import run_llm_jobs3637        async def go() -> None:38            while True:39                stats = await run_llm_jobs(limit=limit)40                out.print(stats)41                if once or not stats["claimed"] or stats.get("skipped_budget"):42                    return4344        run_async(go())4546    @app.command()47    def metrics(all: Annotated[bool, typer.Option("--all", help="every company (default: active in the last 90 days)")] = False,48                company: Annotated[str | None, typer.Option("--company", help="one company slug or id")] = None) -> None:49        """Compute company metrics (activity, hiring momentum, AI adoption, velocity, CCI…) into metrics_current + metric_series."""50        from companyatlas.db import fetch_val, transaction51        from companyatlas.services.metrics import compute_company_metrics5253        async def go() -> None:54            ids = None55            if company:56                async with transaction() as conn:57                    cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company)58                if not cid:59                    console.print(f"[red]company not found: {company}[/]")60                    raise typer.Exit(1)61                ids = [cid]62            out.print(await compute_company_metrics(ids, all_companies=all))63            if company:64                async with transaction() as conn:65                    from companyatlas.db import fetch_all6667                    rows = await fetch_all(conn, "select metric, value, confidence, formula_version, computed_at from metrics_current where company_id = :c order by metric", c=ids[0])68                t = Table(title=f"metrics · {company}")69                for col in ("metric", "value", "confidence", "formula", "computed_at"):70                    t.add_column(col)71                for r in rows:72                    t.add_row(r["metric"], f"{r['value']:.2f}", f"{r['confidence']:.2f}", r["formula_version"], str(r["computed_at"])[:19])73                out.print(t)7475        run_async(go())7677    @app.command()78    def daily(day: Annotated[str | None, typer.Option("--day", help="YYYY-MM-DD (UTC) or 'today'")] = None,79              catch_up: Annotated[bool, typer.Option("--catch-up", help="fill every missing day up to yesterday")] = False,80              include_today: Annotated[bool, typer.Option("--include-today")] = False) -> None:81        """Daily aggregates: company_daily, global_daily (activity index, baseline 100), baselines."""82        from companyatlas.services.metrics import compute_daily, compute_daily_catch_up8384        async def go() -> None:85            if catch_up:86                results = await compute_daily_catch_up(include_today=include_today)87                out.print({"days": len(results), "last": results[-1] if results else None})88                return89            d = datetime.now(UTC).date() if day in (None, "today") else date.fromisoformat(day)90            out.print(await compute_daily(d))9192        run_async(go())9394    @app.command()95    def signals(company: Annotated[str | None, typer.Option("--company")] = None) -> None:96        """Detect company / industry / country signals (hiring surge, launch build-up, expansion, AI acceleration…)."""97        from companyatlas.db import fetch_all, fetch_val, transaction98        from companyatlas.services.signals import compute_signals99100        async def go() -> None:101            ids = None102            if company:103                async with transaction() as conn:104                    cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company)105                ids = [cid] if cid else []106            out.print(await compute_signals(ids))107            async with transaction() as conn:108                rows = await fetch_all(conn, """select coalesce(co.slug, s.scope || ':' || s.scope_key) as who, s.kind, s.strength, s.confidence, s.title109                                                from signals s left join companies co on co.id = s.company_id where s.status = 'active' order by s.strength desc limit 30""")110            t = Table(title="active signals")111            for col in ("scope", "kind", "strength", "confidence", "title"):112                t.add_column(col)113            for r in rows:114                t.add_row(r["who"], r["kind"], f"{r['strength']:.2f}", f"{r['confidence']:.2f}", r["title"])115            out.print(t)116117        run_async(go())118119    @app.command()120    def trends(days: int = 1, window: int = 7) -> None:121        """Extract trending terms from event/news titles for the last N days and print momentum."""122        from companyatlas.services.trends import compute_trends_range, store_momentum_snapshots, trend_momentum123124        async def go() -> None:125            for r in await compute_trends_range(days):126                out.print(r)127            await store_momentum_snapshots()128            items = await trend_momentum(window)129            t = Table(title=f"trend momentum · {window}d")130            for col in ("term", "mentions", "companies", "momentum"):131                t.add_column(col)132            for it in items[:25]:133                t.add_row(it["term"], str(it["mentions"]), str(it["companies"]), f"{it['momentum']:+.2f}")134            out.print(t)135136        run_async(go())137138    @app.command("alerts-eval")139    def alerts_eval(event: Annotated[list[str] | None, typer.Option("--event", help="event id(s); default = sweep")] = None) -> None:140        """Evaluate alerts for given events, or run the catch-up sweep (+ metric alerts)."""141        from companyatlas.services.alerts import evaluate_alerts, sweep142143        run_async(_print(evaluate_alerts(event) if event else sweep()))144145    @app.command("llm-test")146    def llm_test(prompt: str = "Reply with a JSON object {\"ok\": true, \"model\": \"<your model name>\"}", tier: str = "small", json_mode: bool = True) -> None:147        """Live round-trip against the configured LLM endpoint: prints health, model, latency and the answer."""148        from pydantic import BaseModel149150        from companyatlas.config import settings151        from companyatlas.services.llm.gateway import LLMError, get_provider152153        class Probe(BaseModel):154            ok: bool = True155            model: str | None = None156            answer: str | None = None157158        async def go() -> None:159            if not settings.llm_configured:160                console.print("[red]LLM not configured (CA_LLM_BASE_URL / CA_LLM_ENABLED)[/]")161                raise typer.Exit(1)162            p = get_provider()163            h = await p.health()164            out.print({"health": h.ok, "base_url": h.base_url, "latency_ms": h.latency_ms, "models": h.models[:20], "error": h.error})165            t0 = time.monotonic()166            try:167                if json_mode:168                    res = await p.complete_json(tier, "You are a health probe. Answer briefly.", prompt, Probe, max_tokens=120)169                    answer = res.data.model_dump()170                else:171                    res = await p.complete_text(tier, "You are a health probe. Answer briefly.", prompt, max_tokens=120)172                    answer = res.data173            except LLMError as exc:174                console.print(f"[red]LLM error: {exc}[/]")175                raise typer.Exit(1) from exc176            out.print({"model": res.model, "latency_ms": res.latency_ms, "wall_ms": int((time.monotonic() - t0) * 1000), "request_tokens": res.request_tokens,177                       "response_tokens": res.response_tokens, "attempts": res.attempts, "repaired": res.repaired, "answer": answer})178            await p.close()179180        run_async(go())181182    @app.command("reprocess-events")183    def reprocess(since: Annotated[str, typer.Option("--since", help="ISO date/datetime or e.g. 7d")] = "7d",184                  company: Annotated[str | None, typer.Option("--company")] = None, limit: int = 5000) -> None:185        """Re-run deterministic rules on processed changes (no refetch; dedupe keys keep it idempotent)."""186        from companyatlas.db import fetch_val, transaction187        from companyatlas.services.events import reprocess_events188189        async def go() -> None:190            cid = None191            if company:192                async with transaction() as conn:193                    cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company)194            out.print(await reprocess_events(_parse_since(since), limit=limit, company_id=cid))195196        run_async(go())197198    @app.command()199    def events(company: Annotated[str | None, typer.Option("--company")] = None, type: Annotated[str | None, typer.Option("--type")] = None,200               limit: int = 50) -> None:201        """List recent events."""202        from companyatlas.services.events import list_events203204        async def go() -> None:205            rows = await list_events(company=company, event_type=type, limit=limit)206            t = Table(title="events")207            for col in ("detected", "company", "subtype", "imp", "conf", "origin", "status", "title"):208                t.add_column(col)209            for r in rows:210                t.add_row(str(r["detected_at"])[:16], r["slug"], r["event_subtype"], f"{r['importance']:.2f}", r["confidence_label"], r["origin"], r["status"], r["title"][:90])211            out.print(t)212213        run_async(go())214215    @app.command()216    def digest(company: Annotated[str | None, typer.Option("--company")] = None, scope: Annotated[str | None, typer.Option("--scope", help="industry|country")] = None,217               key: Annotated[str | None, typer.Option("--key")] = None, days: int = 7) -> None:218        """Print digest JSON for a company or an industry/country scope."""219        from companyatlas.services.digest import company_digest, scope_digest220221        async def go() -> None:222            if company:223                data = await company_digest(company, days=days)224            elif scope and key:225                data = await scope_digest(scope, key, days=days)226            else:227                console.print("[red]--company or --scope + --key required[/]")228                raise typer.Exit(1)229            out.print_json(json.dumps(data, default=str))230231        run_async(go())232233    @app.command("retention-intel")234    def retention(dry_run: Annotated[bool, typer.Option("--dry-run")] = True) -> None:235        """Prune unchanged observations > 90 d, crawl_runs > 30 d, archive finished llm_jobs > 60 d (dry-run by default)."""236        from companyatlas.services.retention import run_retention237238        run_async(_print(run_retention(dry_run=dry_run)))239240241async def _print(coro) -> None:  # type: ignore[no-untyped-def]242    out.print(await coro)243244245def _parse_since(value: str) -> datetime:246    v = value.strip().lower()247    if v.endswith("d") and v[:-1].isdigit():248        return datetime.now(UTC) - timedelta(days=int(v[:-1]))249    if v.endswith("h") and v[:-1].isdigit():250        return datetime.now(UTC) - timedelta(hours=int(v[:-1]))251    dt = datetime.fromisoformat(value)252    return dt if dt.tzinfo else dt.replace(tzinfo=UTC)253