k8smanager-cli/backend/main.py

"""K8S 智能诊断平台 - FastAPI 主应用"""
import os
import json
import asyncio
from datetime import datetime
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional

from backend.config import PORT
from backend.diagnosis import (
    run_full_diagnosis, build_diagnosis_report, ALL_CHECKS,
    check_cluster_info, check_pods, check_nodes, check_events,
    check_deployments, check_services, check_pvc, check_resource_usage,
    check_network_policies,
)
from backend.ai_agent import analyze_with_ai, chat_with_ai

app = FastAPI(title="K8S 智能诊断平台", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 存储最近诊断结果 (内存)
_last_diagnosis: Optional[dict] = None
_last_report: str = ""

# 诊断项执行顺序 (带中文标题)
CHECK_ORDER = [
    ("cluster_info", "集群信息", check_cluster_info),
    ("nodes", "Node 状态检查", check_nodes),
    ("pods", "Pod 状态检查", check_pods),
    ("deployments", "Deployment 状态检查", check_deployments),
    ("services", "Service/Endpoint 检查", check_services),
    ("events", "异常事件检查", check_events),
    ("pvc", "PVC 存储检查", check_pvc),
    ("resource_usage", "资源使用率", check_resource_usage),
    ("network", "网络策略检查", check_network_policies),
]


class DiagnoseRequest(BaseModel):
    namespace: str = ""
    checks: Optional[list[str]] = None


class AnalyzeRequest(BaseModel):
    question: str = ""


class ChatRequest(BaseModel):
    messages: list[dict]


@app.get("/api/health")
async def health():
    return {"status": "ok", "time": datetime.now().isoformat()}


@app.get("/api/checks")
async def list_checks():
    """列出所有可用诊断项"""
    return {
        "checks": [
            {"name": name, "title": title}
            for name, title, _ in CHECK_ORDER
        ]
    }


@app.get("/api/diagnose/stream")
async def diagnose_stream(namespace: str = ""):
    """SSE 流式诊断 - 实时推送每一步进度"""
    global _last_diagnosis, _last_report

    async def event_generator():
        start = datetime.now()
        results = {}
        all_issues = []
        total = len(CHECK_ORDER)

        for idx, (name, title, fn) in enumerate(CHECK_ORDER, 1):
            # 推送: 开始检查
            yield f"data: {json.dumps({'type': 'progress', 'step': idx, 'total': total, 'name': name, 'title': title, 'status': 'running'}, ensure_ascii=False)}\n\n"

            try:
                if name in ("pods", "events", "deployments", "services", "pvc", "network"):
                    result = await fn(namespace)
                else:
                    result = await fn()
                results[name] = result
                issue_count = len(result.get("issues", []))
                for issue in result.get("issues", []):
                    issue["check"] = title
                    all_issues.append(issue)
            except Exception as e:
                results[name] = {"name": name, "title": title, "error": str(e), "issues": []}
                issue_count = -1

            # 推送: 完成检查
            yield f"data: {json.dumps({'type': 'progress', 'step': idx, 'total': total, 'name': name, 'title': title, 'status': 'done', 'issues': issue_count}, ensure_ascii=False)}\n\n"

        # 汇总
        critical = [i for i in all_issues if i["level"] == "critical"]
        warning = [i for i in all_issues if i["level"] == "warning"]
        elapsed = (datetime.now() - start).total_seconds()
        score = max(0, 100 - len(critical) * 15 - len(warning) * 5)
        status = "healthy" if score >= 90 else ("warning" if score >= 60 else "critical")

        diagnosis = {
            "timestamp": start.isoformat(),
            "elapsed_seconds": round(elapsed, 2),
            "namespace": namespace or "all",
            "score": score,
            "status": status,
            "critical_count": len(critical),
            "warning_count": len(warning),
            "issues": all_issues,
            "checks": results,
        }

        _last_diagnosis = diagnosis
        _last_report = build_diagnosis_report(diagnosis)

        # 推送: 最终结果
        yield f"data: {json.dumps({'type': 'complete', 'diagnosis': diagnosis}, ensure_ascii=False)}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.post("/api/diagnose")
async def diagnose(req: DiagnoseRequest):
    """执行一键诊断 (非流式，兼容旧接口)"""
    global _last_diagnosis, _last_report
    _last_diagnosis = await run_full_diagnosis(namespace=req.namespace, checks=req.checks)
    _last_report = build_diagnosis_report(_last_diagnosis)
    return _last_diagnosis


@app.get("/api/diagnosis/latest")
async def get_latest():
    """获取最近一次诊断结果"""
    if _last_diagnosis is None:
        raise HTTPException(status_code=404, detail="暂无诊断记录，请先执行诊断")
    return _last_diagnosis


@app.post("/api/analyze")
async def analyze(req: AnalyzeRequest):
    """AI 分析诊断结果"""
    global _last_diagnosis, _last_report
    if _last_diagnosis is None:
        _last_diagnosis = await run_full_diagnosis()
        _last_report = build_diagnosis_report(_last_diagnosis)
    result = await analyze_with_ai(_last_report, question=req.question)
    return result


@app.post("/api/chat")
async def chat(req: ChatRequest):
    """AI 多轮对话 (基于诊断上下文追问)"""
    result = await chat_with_ai(req.messages, context=_last_report)
    return result


@app.get("/api/report")
async def get_report():
    """获取诊断报告文本"""
    if not _last_report:
        raise HTTPException(status_code=404, detail="暂无诊断报告")
    return {"report": _last_report}


# 静态文件服务 (前端)
static_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "frontend", "dist")
if os.path.isdir(static_dir):
    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)