atc-agents/api/platform_context.py

"""Live 'platform capabilities + recent changes + masking guidance' block for the LLM.

This is rebuilt on every question from live in-process state, so the assistant is
always autonomously aware of the latest things running in the lab (Data Hub,
Spark Workbench, HDFS→Kafka→Spark→S3 pipeline, autonomous agents) and of the
exact masking policy currently in force.
"""

from __future__ import annotations

from typing import Any


def _fmt_ts(ts: Any) -> str:
    try:
        return str(ts)[:19]
    except Exception:
        return "?"


def build_platform_section() -> str:
    lines: list[str] = ["=== PLATFORM CAPABILITIES & RECENT CHANGES (live) ==="]

    lines += [
        "Command Center features currently deployed:",
        "  - Data Hub: tabbed UI with 'Source Databases' (PostgreSQL, MySQL, MongoDB, Cassandra, Neo4j) and 'Hadoop' (HDFS files, Hive/Iceberg tables, Spark, pipeline).",
        "  - Spark Lakehouse Workbench (Databricks-style): pick any federated table, run preview/filter/aggregate/profile/join/SQL on the distributed engine, and materialize results to Iceberg (S3/HDFS-backed). Live execution matrix: splits, rows, bytes, CPU, wall-time, peak memory, nodes.",
        "  - The same workbench is embedded in each source-database UI (scoped to that source's Trino catalog).",
        "  - Data Flow: live lineage graph with a master pulse switch (Run / Pause / Stop) that starts/stops the animated flow.",
        "  - Pipeline: HDFS (Iceberg historical_sales) → Kafka topic hdfs.historical.sales → Spark transform → Iceberg curated → S3.",
        "Autonomous agents (run continuously, toggleable):",
        "  - Data Custodian (DML): generates live INSERT/UPDATE/DELETE on the source DBs so Debezium CDC streams to Kafka.",
        "  - Data Custodian (Hadoop offload): periodically offloads recent source rows into the Hadoop Iceberg lake (iceberg.hadoop.*_offload), the batch counterpart to CDC.",
        "  - ETL agent: autonomously triggers data movements (HDFS→Iceberg, mask→curated, generators).",
    ]

    # live streaming + flow
    try:
        from streaming_ops import build_streaming_status, flow_snapshot  # type: ignore
        flow = flow_snapshot()
        lines.append(f"Data Flow pulse: {flow.get('mode')}")
    except Exception:
        pass

    # autonomous agent state
    try:
        from agent_ops import _state, _etl_state, _custodian_state  # type: ignore
        lines.append(
            f"DML agent: {'on' if _state.get('enabled') else 'off'} "
            f"(ops_total={_state.get('ops_total')}, last={_state.get('last_op')})"
        )
        lines.append(
            f"ETL agent: {'on' if _etl_state.get('enabled') else 'off'} "
            f"(runs={_etl_state.get('runs_total')}, last={_etl_state.get('last')})"
        )
        cust = _custodian_state
        lines.append(
            f"Custodian Hadoop offload: {'on' if cust.get('enabled') else 'off'} "
            f"(offloads={cust.get('runs_total')}, last={cust.get('last')})"
        )
    except Exception:
        pass

    # recent movements
    try:
        from movements import last_runs  # type: ignore
        runs = last_runs()
        if runs:
            lines.append("Recent data-movement runs:")
            for mid, r in list(runs.items())[-6:]:
                lines.append(f"  - {mid}: {r.get('state')} rows={r.get('rows')} {_fmt_ts(r.get('ended_at'))}")
    except Exception:
        pass

    # recent spark workbench runs
    try:
        from spark_workbench import _runs as wb_runs, _run_order  # type: ignore
        recent = [wb_runs[r] for r in _run_order[-6:] if r in wb_runs]
        if recent:
            lines.append("Recent Spark Workbench runs:")
            for r in recent:
                st = (r.get("stats") or {})
                tgt = f" → {r.get('target')}" if r.get("target") else ""
                lines.append(
                    f"  - {r.get('label')}: {r.get('state')} rows={st.get('processed_rows')}{tgt}"
                )
    except Exception:
        pass

    return "\n".join(lines)


def build_masking_section() -> str:
    """Exact masking policy + strict guidance so the LLM can answer about masked
    data without ever revealing masked raw values."""
    lines: list[str] = ["=== DATA MASKING POLICY (enforced) ==="]
    masked: list[str] = []
    unmasked: list[str] = []
    try:
        from pii_catalog import get_pii  # type: ignore
        data = get_pii()
        for d in data.get("datasets", []):
            for c in d.get("pii_columns", []):
                tag = f"{d.get('label')}.{c.get('name')} [{c.get('category')}]"
                (masked if c.get("masked") else unmasked).append(tag)
        summ = data.get("summary", {})
        lines.append(
            f"PII columns: {summ.get('pii_columns', 0)} total — "
            f"{summ.get('masked_columns', 0)} masked, {summ.get('unmasked_columns', 0)} visible."
        )
    except Exception as exc:
        lines.append(f"(masking catalog unavailable: {exc})")

    if masked:
        lines.append("MASKED columns (raw values are withheld — token 🔒 MASKED):")
        for m in masked[:40]:
            lines.append(f"  - {m}")
    if unmasked:
        lines.append("Visible PII columns (operator opted out of masking):")
        for u in unmasked[:40]:
            lines.append(f"  - {u}")

    lines += [
        "",
        "How to handle masked data when answering:",
        "  1. NEVER reveal, guess, reconstruct or print the raw value of a MASKED column. If a value comes in as '🔒 MASKED', keep it masked.",
        "  2. DO still answer helpfully: confirm the column exists and is masked for privacy/governance, and explain why (PII protection policy).",
        "  3. You MAY use and report non-sensitive aggregates, counts, distributions and derived metrics over masked columns (e.g. 'there are N distinct customers') as long as no individual raw value is exposed.",
        "  4. Tell the operator they can unmask a specific column from the Data Flow PII overlay if they have the authority, and that the curated/masked Iceberg layer is physically masked and cannot be unmasked.",
        "  5. Unmasked PII columns may be shown, but flag that they are sensitive.",
    ]
    return "\n".join(lines)


def build_business_data_section() -> str:
    """Detailed inventory of the actual business data — tables, columns, types,
    row counts and per-column masked/visible status — so the assistant knows the
    data in detail and exactly what is masked vs not."""
    lines: list[str] = ["=== BUSINESS DATA CATALOG (live, every detail) ==="]

    # source row totals (cheap estimates)
    try:
        import sql_console as s
        totals = {
            "PostgreSQL sales_orders": s._table_row_count("postgres", "public.sales_orders"),
            "MySQL employee_events": s._table_row_count("mysql", "hr.employee_events"),
            "MongoDB supplychain.events": s._table_row_count("mongodb", "supplychain.events"),
        }
        lines.append("Source volumes (live row counts):")
        for k, v in totals.items():
            lines.append(f"  - {k}: {v:,} rows" if isinstance(v, int) else f"  - {k}: ~")
    except Exception:
        pass

    # full column dictionary + masking per column
    try:
        from trino_federated import build_dictionary
        d = build_dictionary()
        summ = d.get("summary", {})
        lines.append(
            f"Catalogued tables: {summ.get('tables', 0)} — columns: {summ.get('columns', 0)}, "
            f"PII columns: {summ.get('pii_columns', 0)} ({summ.get('masked_columns', 0)} masked)."
        )
        for t in d.get("tables", []):
            lines.append(f"\n{t['engine']} · {t['fqn']} — {t['desc']}")
            for c in t.get("columns", []):
                flag = ""
                if c.get("masked"):
                    flag = f"  [MASKED · {c.get('category')}]"
                elif c.get("pii"):
                    flag = f"  [PII visible · {c.get('category')}]"
                lines.append(f"    · {c['name']} ({c['type']}){flag}")
    except Exception as exc:
        lines.append(f"(data dictionary unavailable: {exc})")

    # federated cross-source marquee (region scorecard), if computed
    try:
        from trino_federated import _marquee, _load_marquee
        if _marquee.get("data") is None:
            _load_marquee()
        m = _marquee.get("data")
        if m and m.get("rows"):
            lines.append("\nFederated region scorecard (one Trino SQL across PostgreSQL+MySQL+MongoDB):")
            for r in m["rows"][:8]:
                rev = r.get("revenue")
                rev_s = f"€{rev:,.0f}" if isinstance(rev, (int, float)) else "?"
                lines.append(
                    f"  - {r.get('region')}: orders={r.get('orders')}, revenue={rev_s}, "
                    f"hr_events={r.get('hr_events')}, supply_events={r.get('supply_events')}"
                )
    except Exception:
        pass

    lines.append(
        "\nData is also materialized into the Hadoop lake as external Iceberg tables "
        "(iceberg.hadoop.*_ext) and exposed through one federated Trino engine "
        "(catalogs: postgres_sales, mysql_hr, mongodb_supplychain, cassandra_telemetry, iceberg, kafka)."
    )
    return "\n".join(lines)


def build_llm_addendum() -> str:
    try:
        platform = build_platform_section()
    except Exception as exc:
        platform = f"(platform section error: {exc})"
    try:
        business = build_business_data_section()
    except Exception as exc:
        business = f"(business data section error: {exc})"
    try:
        masking = build_masking_section()
    except Exception as exc:
        masking = f"(masking section error: {exc})"
    return "\n\n".join([platform, business, masking])