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_llm_addendum() -> str:
    try:
        platform = build_platform_section()
    except Exception as exc:
        platform = f"(platform section error: {exc})"
    try:
        masking = build_masking_section()
    except Exception as exc:
        masking = f"(masking section error: {exc})"
    return "\n\n".join([platform, masking])