spb/company-atlas

1"""`catlas` profile enrichment commands: `enrich-companies`, `profile` (services/enrichment.py)."""2from __future__ import annotations34import json5from typing import Annotated67import typer8from rich.table import Table91011def register(app: typer.Typer) -> None:12    @app.command("enrich-companies")13    def enrich_companies(14        limit: Annotated[int | None, typer.Option("--limit", help="companies per run (default: CA_ENRICH_BATCH)")] = None,15        company: Annotated[list[str] | None, typer.Option("--company", help="slug, id or QID (repeatable); forces re-enrichment")] = None,16        source: Annotated[str, typer.Option("--source", help="wikidata|wikipedia|homepage|llm|all (comma-separated allowed)")] = "all",17        concurrency: Annotated[int | None, typer.Option("--concurrency")] = None,18        no_llm: Annotated[bool, typer.Option("--no-llm", help="never call the LLM even when configured")] = False,19    ) -> None:20        """Enrich company profiles (Wikidata → Wikipedia → homepage → optional LLM) with provenance; never-enriched companies first."""21        from companyatlas.cli import console, out, run_async22        from companyatlas.services.enrichment import SOURCES, enrich_pending2324        wanted = tuple(SOURCES) if source.strip().lower() in ("all", "") else tuple(s.strip().lower() for s in source.split(",") if s.strip())25        unknown = [s for s in wanted if s not in SOURCES]26        if unknown:27            console.print(f"[red]unknown source(s): {', '.join(unknown)} — choose from {', '.join(SOURCES)} or all[/]")28            raise typer.Exit(1)29        stats = run_async(enrich_pending(limit=limit, concurrency=concurrency, sources=wanted, llm=not no_llm, company_keys=company))30        table = Table(title="enrich-companies", show_header=False)31        for k, v in stats.items():32            table.add_row(k, json.dumps(v) if isinstance(v, dict) else str(v))33        out.print(table)3435    @app.command("profile")36    def profile_cmd(key: Annotated[str, typer.Argument(help="company slug, id or Wikidata QID")],37                    as_json: Annotated[bool, typer.Option("--json", help="raw profile JSON")] = False) -> None:38        """Print a company's stored profile (source_meta.profile) with the key facts and their sources."""39        from companyatlas.cli import console, out, run_async40        from companyatlas.db import connection, fetch_all41        from companyatlas.services.enrichment import load_company, person_source, profile_facts4243        async def go() -> tuple[dict | None, list[dict], list[dict]]:44            async with connection() as conn:45                c = await load_company(conn, key)46                if c is None:47                    return None, [], []48                people = await fetch_all(conn, "select name, title, status, source_url from people where company_id = :c order by is_executive desc, name limit 40", c=c["id"])49                rels = await fetch_all(conn, "select r.kind, coalesce(o.display_name, r.to_name) as name, r.valid_from, r.valid_to, r.provenance from company_relationships r "50                                             "left join companies o on o.id = r.to_company_id where r.from_company_id = :c order by r.kind, name limit 60", c=c["id"])51                return c, people, rels5253        company, people, rels = run_async(go())54        if company is None:55            console.print(f"[red]company not found: {key}[/]")56            raise typer.Exit(1)57        meta = company.get("source_meta") or {}58        if isinstance(meta, str):59            meta = json.loads(meta)60        profile = meta.get("profile")61        if not profile:62            console.print(f"[yellow]{company['slug']} has not been enriched yet — run `catlas enrich-companies --company {company['slug']}`[/]")63            raise typer.Exit(2)64        if as_json:65            out.print_json(json.dumps(profile, default=str, ensure_ascii=False))66            return67        out.print(f"[bold]{company['display_name']}[/] · {company['canonical_domain']} · enriched {profile.get('enriched_at')} · sources: "68                  f"{', '.join(meta.get('enrichment', {}).get('sources') or [])}")69        if profile.get("description"):70            out.print(f"\n{profile['description']}\n[dim]— {profile.get('description_source')} · {profile.get('description_attribution') or ''} "71                      f"{profile.get('description_url') or ''}[/]\n")72        t = Table(title="facts")73        for col in ("fact", "value", "source", "url"):74            t.add_column(col)75        for f in profile_facts(profile):76            t.add_row(f["label"], str(f["value"]), f.get("source") or "", (f.get("url") or "")[:80])77        out.print(t)78        if profile.get("industries") or profile.get("products"):79            out.print(f"industries: {', '.join(profile.get('industries') or [])}  ·  labels: {', '.join(profile.get('industry_labels') or [])}")80            out.print(f"products: {', '.join(profile.get('products') or [])}")81        if profile.get("socials"):82            out.print("socials: " + "  ".join(f"{k}={v}" for k, v in profile["socials"].items()))83        if people:84            p = Table(title="people")85            for col in ("name", "title", "status", "source"):86                p.add_column(col)87            for r in people:88                p.add_row(r["name"], r.get("title") or "", r["status"], person_source(r.get("source_url")))89            out.print(p)90        if rels:91            r_ = Table(title="relationships")92            for col in ("kind", "company", "valid_from", "valid_to", "property"):93                r_.add_column(col)94            for r in rels:95                prov = r.get("provenance") or {}96                if isinstance(prov, str):97                    prov = json.loads(prov)98                r_.add_row(r["kind"], r.get("name") or "", str(r.get("valid_from") or ""), str(r.get("valid_to") or ""), prov.get("property") or "")99            out.print(r_)100