atc-agents/main.py

"""ATC Command Center API — FastAPI backend."""

from __future__ import annotations

import asyncio
import json
import os
import time
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Any

import httpx
import redis.asyncio as aioredis
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from agent_terminal import (
    get_all_terminals,
    get_terminal_lines,
    init_terminals,
    make_logger,
    set_terminal_publisher,
    terminal_log,
)
from lab_context import collect_full_lab_context, format_context_for_agent
from workload import build_workload_payload
from pydantic import BaseModel, Field
from sqlalchemy import Column, DateTime, String, Text, create_engine, select
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker

REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
DOCKHAND_URL = os.getenv("DOCKHAND_URL", "http://10.0.21.45:8082")
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:////data/atc-agents.db")
GPU_URL = os.getenv("GPU_URL", "http://10.0.20.106:9000")
GPU_UI_URL = os.getenv("GPU_UI_URL", GPU_URL)
LLM_URL = os.getenv("LLM_URL", "http://10.0.20.106:8001/v1")
LLM_MODEL = os.getenv("LLM_MODEL", "qwen2.5-32b-gptq")
LLM_API_KEY = os.getenv("LLM_API_KEY", "sk-local")
LLM_TIMEOUT = float(os.getenv("LLM_TIMEOUT", "120"))

AGENTS = [
    {
        "id": "etl-guardian",
        "name": "ETL Guardian",
        "color": "#00f0ff",
        "zone": "etl",
        "role": "Airflow, Kafka, Debezium, S3 pipeline",
        "icon": "⚡",
        "motto": "Pipelines never sleep",
        "capabilities": ["Airflow", "Kafka", "Debezium", "S3", "Connectors"],
        "suggested_prompts": [
            "Hoe staat Debezium er voor?",
            "Zijn alle Airflow DAGs healthy?",
            "Kafka connector status?",
        ],
    },
    {
        "id": "lakehouse-ops",
        "name": "Lakehouse Ops",
        "color": "#ff00aa",
        "zone": "lakehouse",
        "role": "Spark, Trino, Iceberg",
        "icon": "🏔️",
        "motto": "Query the lake, trust the table",
        "capabilities": ["Spark", "Trino", "Iceberg", "Delta", "SQL"],
        "suggested_prompts": [
            "Lakehouse stack status?",
            "Is Trino bereikbaar?",
            "Hoeveel lakehouse containers draaien?",
        ],
    },
    {
        "id": "data-custodian",
        "name": "Data Custodian",
        "color": "#ffaa00",
        "zone": "db",
        "role": "PostgreSQL, MySQL, Mongo, Cassandra, Neo4j",
        "icon": "🛡️",
        "motto": "Guardian of every row",
        "capabilities": ["PostgreSQL", "MySQL", "MongoDB", "Cassandra", "Neo4j"],
        "suggested_prompts": [
            "Database containers status?",
            "Welke DB's draaien niet?",
            "Postgres health check",
        ],
    },
    {
        "id": "hadoop-ranger",
        "name": "Hadoop Ranger",
        "color": "#39ff14",
        "zone": "hadoop",
        "role": "HDFS, YARN cluster",
        "icon": "🌲",
        "motto": "Patrol the data forest",
        "capabilities": ["HDFS", "YARN", "NameNode", "DataNodes"],
        "suggested_prompts": [
            "Is HDFS NameNode up?",
            "Hadoop cluster status?",
            "YARN nodes healthy?",
        ],
    },
    {
        "id": "infra-sentinel",
        "name": "Infra Sentinel",
        "color": "#b366ff",
        "zone": "docker",
        "role": "Docker, Proxmox, GPU, monitoring",
        "icon": "👁️",
        "motto": "See everything, miss nothing",
        "capabilities": ["Docker", "Proxmox", "GPU", "vLLM", "Monitoring"],
        "suggested_prompts": [
            "Hoe staat de GPU?",
            "Welk LLM model draait er?",
            "Docker containers overzicht",
        ],
    },
]

ZONES = [
    {"id": "docker", "label": "DOCKER RACK", "x": 8, "color": "#b366ff"},
    {"id": "db", "label": "DB VAULT", "x": 28, "color": "#ffaa00"},
    {"id": "lakehouse", "label": "LAKEHOUSE HUB", "x": 50, "color": "#ff00aa"},
    {"id": "hadoop", "label": "HADOOP CLUSTER", "x": 72, "color": "#39ff14"},
    {"id": "etl", "label": "ETL PIPE", "x": 92, "color": "#00f0ff"},
]

INTENT_KEYWORDS: dict[str, list[str]] = {
    "data-custodian": ["database", "postgres", "postgresql", "mysql", "mongo", "mongodb", "cassandra", "neo4j", "sql", "db "],
    "lakehouse-ops": ["trino", "spark", "lakehouse", "iceberg", "query", "table"],
    "hadoop-ranger": [
        "hadoop", "hdfs", "yarn", "datanode", "namenode", "replicatie", "replication",
        "rf factor", "opslag", "bestanden", "blocks", "cluster opslag", "data op",
    ],
    "infra-sentinel": ["docker", "container", "vm", "proxmox", "infra", "grafana", "gpu", "vllm", "llm", "nvidia", "inference", "model"],
    "etl-guardian": ["airflow", "dag", "debezium", "kafka", "connector", "etl", "pipeline", "s3"],
}


class Base(DeclarativeBase):
    pass


class FeedEntry(Base):
    __tablename__ = "feed"

    id = Column(String, primary_key=True)
    ts = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    agent_id = Column(String)
    level = Column(String, default="info")
    message = Column(Text)


class Approval(Base):
    __tablename__ = "approvals"

    id = Column(String, primary_key=True)
    ts = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    agent_id = Column(String)
    action = Column(Text)
    reason = Column(Text)
    status = Column(String, default="pending")


engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine)
Base.metadata.create_all(engine)

redis_client: aioredis.Redis | None = None
ws_clients: set[WebSocket] = set()


class PromptRequest(BaseModel):
    message: str = Field(min_length=1, max_length=2000)
    agent_id: str | None = None


class ApprovalDecision(BaseModel):
    approved: bool


def route_agent(message: str) -> str:
    lower = message.lower()
    # Storage/data questions default to Hadoop unless clearly about databases
    if any(w in lower for w in ("data", "opslag", "gb", "replicatie", "replication", "hdfs", "hadoop")):
        if not any(w in lower for w in ("postgres", "mysql", "mongo", "database", "sql", "neo4j", "cassandra")):
            return "hadoop-ranger"
    scores = {aid: sum(1 for kw in kws if kw in lower) for aid, kws in INTENT_KEYWORDS.items()}
    best = max(scores, key=scores.get)
    if scores[best] == 0:
        return "infra-sentinel"
    return best


async def gather_agent_context(
    agent_id: str,
    status: dict[str, Any],
    log: Any | None = None,
) -> str:
    """Full lab snapshot for vLLM — all domains, agent's primary domain highlighted."""
    snapshot = await collect_full_lab_context(gpu_data=status.get("gpu"), log=log)
    snapshot["domains_summary"] = status.get("domains", {})
    ctx = format_context_for_agent(agent_id, snapshot)
    if log:
        await log("info", "fetch", f"▸ Context assembled: {len(ctx)} chars for LLM")
    return ctx


async def ask_llm(
    agent_id: str,
    message: str,
    context: str,
    log: Any | None = None,
) -> str | None:
    agent = next(a for a in AGENTS if a["id"] == agent_id)
    system = f"""Je bent {agent['name']}, een autonomous ops agent in het Dell ATC data lab.
Specialisatie: {agent['role']}.
Motto: {agent.get('motto', '')}

Je antwoordt namens je domein maar hebt zicht op de HELE lab stack: Docker, databases, lakehouse (Trino/Spark/Kafka Connect), ETL (Airflow/Kafka), Hadoop HDFS, en GPU/vLLM.

Regels:
- Antwoord in dezelfde taal als de gebruiker (Nederlands of Engels).
- Gebruik ALLEEN de live data hieronder — verzin geen hosts, poorten, cijfers of connector namen.
- Gebruik exact de container/connector namen uit de data (bijv. mysql-hr-connector, niet "Debezium").
- Als iets DOWN of 0 GB is, zeg dat eerlijk.
- Kort en behulpzaam (max ~10 zinnen); bullet lists mogen als het overzicht helpt.

--- LIVE LAB DATA (primary domain eerst, daarna volledige stack) ---
{context}
"""
    if log:
        await log("info", "llm", f"▸ Querying vLLM model={LLM_MODEL}")
        await log("cmd", "llm", f"$ POST {LLM_URL.rstrip('/')}/chat/completions")
        await log("info", "llm", f"  user: {message[:160]}{'…' if len(message) > 160 else ''}")
    try:
        async with httpx.AsyncClient(timeout=LLM_TIMEOUT) as client:
            t0 = time.monotonic()
            r = await client.post(
                f"{LLM_URL.rstrip('/')}/chat/completions",
                headers={
                    "Authorization": f"Bearer {LLM_API_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": LLM_MODEL,
                    "messages": [
                        {"role": "system", "content": system},
                        {"role": "user", "content": message},
                    ],
                    "max_tokens": 800,
                    "temperature": 0.25,
                },
            )
            r.raise_for_status()
            content = r.json()["choices"][0]["message"]["content"].strip()
            ms = int((time.monotonic() - t0) * 1000)
            if content and content.strip("!"):
                if log:
                    await log("ok", "llm", f"← vLLM response {len(content)} chars ({ms}ms)")
                    preview = content.replace("\n", " ")[:180]
                    await log("info", "llm", f"  » {preview}{'…' if len(content) > 180 else ''}")
                return content
            if log:
                await log("warn", "llm", f"← Empty or invalid LLM output ({ms}ms)")
    except Exception as exc:
        if log:
            await log("err", "llm", f"✗ vLLM error: {exc}")
    return None


def fallback_answer(agent_id: str, context: str) -> str:
    agent_name = next(a["name"] for a in AGENTS if a["id"] == agent_id)
    return f"**{agent_name}** (offline LLM — ruwe data):\n\n{context}"


async def publish_event(event: dict[str, Any]) -> None:
    payload = json.dumps(event, default=str)
    if redis_client:
        await redis_client.publish("ops", payload)
    dead = []
    for ws in ws_clients:
        try:
            await ws.send_text(payload)
        except Exception:
            dead.append(ws)
    for ws in dead:
        ws_clients.discard(ws)


def add_feed(agent_id: str, message: str, level: str = "info") -> dict:
    entry_id = str(uuid.uuid4())[:8]
    with SessionLocal() as db:
        row = FeedEntry(id=entry_id, agent_id=agent_id, message=message, level=level)
        db.add(row)
        db.commit()
    return {
        "id": entry_id,
        "ts": datetime.now(timezone.utc).isoformat(),
        "agent_id": agent_id,
        "message": message,
        "level": level,
    }


async def dockhand_env_containers(env_id: int) -> list[dict]:
    try:
        async with httpx.AsyncClient(timeout=8.0) as client:
            r = await client.get(f"{DOCKHAND_URL}/api/containers", params={"env": env_id})
            r.raise_for_status()
            return r.json()
    except Exception:
        return []


async def probe_url(url: str) -> bool:
    try:
        async with httpx.AsyncClient(timeout=4.0, verify=False) as client:
            r = await client.get(url)
            return r.status_code < 500
    except Exception:
        return False


async def collect_gpu() -> dict[str, Any]:
    host = GPU_URL.replace("http://", "").replace("https://", "").split("/")[0]
    base = {"ok": False, "host": host, "ui_url": GPU_UI_URL}

    try:
        async with httpx.AsyncClient(timeout=6.0) as client:
            metrics_r, model_r, integration_r = await asyncio.gather(
                client.get(f"{GPU_URL}/api/gpu/metrics"),
                client.get(f"{GPU_URL}/api/active-model"),
                client.get(f"{GPU_URL}/api/integration"),
                return_exceptions=True,
            )

            gpus: list[dict[str, Any]] = []
            if isinstance(metrics_r, httpx.Response) and metrics_r.status_code == 200:
                current = metrics_r.json().get("current", {})
                gpus = [
                    {
                        "index": g["index"],
                        "name": g["name"],
                        "util_gpu": g.get("util_gpu", 0),
                        "memory_used_mib": g.get("memory_used_mib", 0),
                        "memory_total_mib": g.get("memory_total_mib", 0),
                        "temperature_c": g.get("temperature_c", 0),
                        "power_w": g.get("power_w", 0),
                    }
                    for g in current.get("gpus", [])
                ]

            active_model = None
            inference_active = False
            vllm_url = None

            if isinstance(model_r, httpx.Response) and model_r.status_code == 200:
                model_data = model_r.json()
                active_model = model_data.get("name")
                inference_active = bool(model_data.get("inference_active"))
                vllm_url = model_data.get("base_url")

            if isinstance(integration_r, httpx.Response) and integration_r.status_code == 200:
                integ = integration_r.json()
                if not active_model:
                    active_model = integ.get("active_name")
                if not inference_active:
                    inference_active = bool(integ.get("inference_active"))
                if not vllm_url:
                    vllm_url = integ.get("recommended_base_url")

            return {
                **base,
                "ok": len(gpus) > 0 or inference_active,
                "inference_active": inference_active,
                "active_model": active_model,
                "vllm_url": vllm_url,
                "gpu_count": len(gpus),
                "gpus": gpus,
            }
    except Exception as exc:
        return {**base, "error": str(exc)}


async def collect_status() -> dict[str, Any]:
    db_containers = await dockhand_env_containers(5)
    db_running = sum(1 for c in db_containers if c.get("state") == "running")
    db_total = len(db_containers) or 6

    lake_containers = await dockhand_env_containers(9)
    lake_running = sum(1 for c in lake_containers if c.get("state") == "running")
    lake_total = len(lake_containers) or 6

    docker_containers = await dockhand_env_containers(1)
    docker_running = sum(1 for c in docker_containers if c.get("state") == "running")

    hdfs_ok = await probe_url("http://10.0.21.61:9870")
    kafka_ok = await probe_url("http://10.0.21.36:9000")
    airflow_ok = await probe_url("http://10.0.21.55:8080")

    def level(running: int, total: int) -> str:
        if total == 0:
            return "unknown"
        ratio = running / total
        if ratio >= 0.9:
            return "ok"
        if ratio >= 0.5:
            return "warn"
        return "down"

    gpu = await collect_gpu()
    gpu_level = "ok" if gpu.get("ok") and gpu.get("inference_active") else ("warn" if gpu.get("ok") else "down")
    gpu_label = gpu.get("active_model") or (f"{gpu.get('gpu_count', 0)} GPUs" if gpu.get("ok") else "offline")

    return {
        "ts": datetime.now(timezone.utc).isoformat(),
        "domains": {
            "docker": {"level": "ok" if docker_running >= 5 else "warn", "label": f"{docker_running} containers", "running": docker_running},
            "databases": {"level": level(db_running, db_total), "label": f"{db_running}/{db_total} up", "running": db_running, "total": db_total},
            "lakehouse": {"level": level(lake_running, lake_total), "label": f"{lake_running}/{lake_total} up", "running": lake_running, "total": lake_total},
            "hadoop": {"level": "ok" if hdfs_ok else "warn", "label": "NN up" if hdfs_ok else "NN check"},
            "etl": {"level": "ok" if kafka_ok and airflow_ok else "warn", "label": "Kafka+Airflow"},
            "gpu": {"level": gpu_level, "label": gpu_label},
        },
        "gpu": gpu,
        "kafka_ok": kafka_ok,
        "airflow_ok": airflow_ok,
        "hdfs_ok": hdfs_ok,
    }


async def run_agent_task(agent_id: str, message: str, prompt_id: str) -> str:
    zone = next(a["zone"] for a in AGENTS if a["id"] == agent_id)
    agent_name = next(a["name"] for a in AGENTS if a["id"] == agent_id)
    log = make_logger(agent_id, prompt_id)

    await terminal_log(
        agent_id,
        f"[{prompt_id}] ▶ Mission accepted: {message}",
        level="info",
        phase="dispatch",
        prompt_id=prompt_id,
    )
    await publish_event({"type": "agent_dispatch", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})
    await asyncio.sleep(0.4)
    await terminal_log(agent_id, f"[{prompt_id}] Walking to zone: {zone}", level="info", phase="dispatch", prompt_id=prompt_id)
    await publish_event({"type": "agent_fetch", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})

    await log("info", "fetch", f"[{prompt_id}] Collecting live lab metrics…")
    status = await collect_status()
    context = await gather_agent_context(agent_id, status, log=log)

    answer = await ask_llm(agent_id, message, context, log=log)
    if not answer:
        await log("warn", "llm", "LLM fallback — returning raw context")
        answer = fallback_answer(agent_id, context)

    await asyncio.sleep(0.3)
    await terminal_log(agent_id, f"[{prompt_id}] ✓ Mission complete", level="ok", phase="done", prompt_id=prompt_id)
    await publish_event({"type": "agent_return", "agent_id": agent_id, "zone": zone, "prompt_id": prompt_id})
    feed = add_feed(agent_id, f"{agent_name} answered: {message[:60]}", "info")
    await publish_event({"type": "feed", "entry": feed})
    await publish_event({"type": "prompt_result", "prompt_id": prompt_id, "agent_id": agent_id, "answer": answer})
    return answer


async def heartbeat_loop() -> None:
    while True:
        try:
            status = await collect_status()
            workload = await collect_workload()
            await publish_event({"type": "status", "data": status})
            await publish_event({"type": "workload", "data": workload})
            for domain, info in status["domains"].items():
                if info["level"] == "down":
                    agent = "data-custodian" if domain == "databases" else "infra-sentinel"
                    feed = add_feed(agent, f"Alert: {domain} is DOWN ({info['label']})", "warn")
                    await publish_event({"type": "feed", "entry": feed})
        except Exception as exc:
            await publish_event({"type": "error", "message": str(exc)})
        await asyncio.sleep(60)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global redis_client
    redis_client = aioredis.from_url(REDIS_URL, decode_responses=True)
    set_terminal_publisher(publish_event)
    init_terminals([a["id"] for a in AGENTS])
    for a in AGENTS:
        await terminal_log(a["id"], f"{a['name']} terminal online — awaiting missions", level="info", phase="boot")
    task = asyncio.create_task(heartbeat_loop())
    add_feed("infra-sentinel", "ATC Command Center API online", "info")
    yield
    task.cancel()
    if redis_client:
        await redis_client.close()


app = FastAPI(title="ATC Command Center API", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api/health")
async def health():
    llm_ok = False
    try:
        async with httpx.AsyncClient(timeout=4.0) as client:
            r = await client.get(f"{LLM_URL.rstrip('/')}/models", headers={"Authorization": f"Bearer {LLM_API_KEY}"})
            llm_ok = r.status_code == 200
    except Exception:
        pass
    return {"ok": True, "ts": datetime.now(timezone.utc).isoformat(), "llm_url": LLM_URL, "llm_ok": llm_ok, "llm_model": LLM_MODEL}


async def collect_workload() -> dict[str, Any]:
    gpu = await collect_gpu()
    snap = await collect_full_lab_context(gpu_data=gpu)
    return build_workload_payload(snap)


@app.get("/api/workload")
async def get_workload():
    return await collect_workload()


@app.get("/api/status")
async def get_status():
    return await collect_status()


@app.get("/api/gpu")
async def get_gpu():
    return await collect_gpu()


def agent_stats() -> dict[str, dict[str, Any]]:
    stats: dict[str, dict[str, Any]] = {a["id"]: {"tasks": 0, "last_active": None, "alerts": 0} for a in AGENTS}
    with SessionLocal() as db:
        rows = db.execute(select(FeedEntry).order_by(FeedEntry.ts.desc()).limit(200)).scalars().all()
    for r in rows:
        aid = r.agent_id
        if aid not in stats:
            continue
        stats[aid]["tasks"] += 1
        if r.level == "warn":
            stats[aid]["alerts"] += 1
        if stats[aid]["last_active"] is None and r.ts:
            stats[aid]["last_active"] = r.ts.isoformat()
    return stats


@app.get("/api/agents")
async def get_agents():
    stats = agent_stats()
    enriched = [{**a, "stats": stats.get(a["id"], {})} for a in AGENTS]
    return {"agents": enriched, "zones": ZONES}


@app.get("/api/terminals")
async def get_terminals(limit: int = 200):
    return {"terminals": get_all_terminals(limit)}


@app.get("/api/terminals/{agent_id}")
async def get_agent_terminal(agent_id: str, limit: int = 200):
    valid = {a["id"] for a in AGENTS}
    if agent_id not in valid:
        return {"error": "unknown agent"}
    return {"agent_id": agent_id, "lines": get_terminal_lines(agent_id, limit)}


@app.get("/api/feed")
async def get_feed(limit: int = 50):
    with SessionLocal() as db:
        rows = db.execute(select(FeedEntry).order_by(FeedEntry.ts.desc()).limit(limit)).scalars().all()
    return {
        "entries": [
            {
                "id": r.id,
                "ts": r.ts.isoformat() if r.ts else None,
                "agent_id": r.agent_id,
                "message": r.message,
                "level": r.level,
            }
            for r in rows
        ]
    }


@app.get("/api/approvals")
async def get_approvals():
    with SessionLocal() as db:
        rows = db.execute(select(Approval).where(Approval.status == "pending")).scalars().all()
    return {
        "approvals": [
            {
                "id": r.id,
                "ts": r.ts.isoformat() if r.ts else None,
                "agent_id": r.agent_id,
                "action": r.action,
                "reason": r.reason,
                "status": r.status,
            }
            for r in rows
        ]
    }


@app.post("/api/approvals/{approval_id}/decide")
async def decide_approval(approval_id: str, body: ApprovalDecision):
    with SessionLocal() as db:
        row = db.get(Approval, approval_id)
        if not row:
            return {"error": "not found"}
        row.status = "approved" if body.approved else "denied"
        db.commit()
        agent_id = row.agent_id
        action = row.action
    msg = f"Approval {'approved' if body.approved else 'denied'}: {action}"
    feed = add_feed(agent_id, msg, "info" if body.approved else "warn")
    await publish_event({"type": "feed", "entry": feed})
    await publish_event({"type": "approval_update", "id": approval_id, "status": row.status})
    return {"ok": True, "status": row.status}


@app.post("/api/prompt")
async def post_prompt(body: PromptRequest):
    prompt_id = str(uuid.uuid4())[:8]
    valid_ids = {a["id"] for a in AGENTS}
    agent_id = body.agent_id if body.agent_id in valid_ids else route_agent(body.message)
    add_feed(agent_id, f"Prompt received: {body.message}", "info")
    asyncio.create_task(run_agent_task(agent_id, body.message, prompt_id))
    return {"prompt_id": prompt_id, "agent_id": agent_id, "status": "dispatched"}


@app.websocket("/api/ws/ops")
async def ws_ops(websocket: WebSocket):
    await websocket.accept()
    ws_clients.add(websocket)
    try:
        status = await collect_status()
        workload = await collect_workload()
        await websocket.send_text(json.dumps({"type": "status", "data": status}, default=str))
        await websocket.send_text(json.dumps({"type": "workload", "data": workload}, default=str))
        await websocket.send_text(json.dumps({
            "type": "terminal_history",
            "terminals": get_all_terminals(150),
        }, default=str))
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        ws_clients.discard(websocket)