web_server.py

"""
简单的HTTP服务器，用于实时监控
"""
import http.server
import socketserver
import json
import re
from pathlib import Path
from urllib.parse import urlparse, parse_qs

PORT = 8000
SIM_ROOTS = [Path("data/logs"), Path("logs")]


def _agent_memory_dir_name(agent_id: str) -> str:
    """与 core.agent.Agent.parse_slot_id 一致：磁盘目录为 agent_{slot}/。"""
    text = str(agent_id or "")
    m_new = re.match(r"^agent_(\d+)_gen\d+$", text)
    if m_new:
        return f"agent_{m_new.group(1)}"
    m_old = re.match(r"^agent_gen\d+_(\d+)$", text)
    if m_old:
        return f"agent_{m_old.group(1)}"
    m_any = re.search(r"(\d+)", text)
    return f"agent_{m_any.group(1)}" if m_any else text


def list_simulation_dirs():
    sims = []
    for root in SIM_ROOTS:
        if root.exists():
            sims.extend([p for p in root.iterdir() if p.is_dir()])
    sims = list({str(p): p for p in sims}.values())
    sims.sort(key=lambda p: p.name, reverse=True)
    return sims


def resolve_simulation_dir(simulation_id: str = "") -> Path:
    sims = list_simulation_dirs()
    if simulation_id:
        for root in SIM_ROOTS:
            target = root / simulation_id
            if target.exists() and target.is_dir():
                return target
    if sims:
        return sims[0]
    return SIM_ROOTS[0] / "default"

class RealTimeHandler(http.server.SimpleHTTPRequestHandler):
    def do_GET(self):
        parsed = urlparse(self.path)
        params = parse_qs(parsed.query)

        if parsed.path == '/api/leaderboard':
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()

            simulation_id = params.get('simulation', [''])[0]
            simulation_dir = resolve_simulation_dir(simulation_id)
            agents_file = simulation_dir / "agents.jsonl"
            realtime_file = simulation_dir / "realtime_state.json"

            env_payload = {}
            elim_from_state = []
            if realtime_file.exists():
                try:
                    with open(realtime_file, 'r', encoding='utf-8') as f:
                        rt = json.load(f)
                    env_payload = (rt.get("environment") or {})
                    elim_from_state = rt.get("eliminated_agents") or []
                except Exception:
                    pass

            if not agents_file.exists():
                self.wfile.write(json.dumps({
                    "simulation": simulation_dir.name,
                    "leaderboard": [],
                    "eliminated": [],
                    "environment": env_payload,
                    "round": 0,
                }, ensure_ascii=False).encode('utf-8'))
                return

            try:
                with open(agents_file, 'r', encoding='utf-8') as f:
                    lines = [line.strip() for line in f.readlines() if line.strip()]
                if not lines:
                    self.wfile.write(json.dumps({
                        "simulation": simulation_dir.name,
                        "leaderboard": [],
                        "eliminated": [],
                        "environment": env_payload,
                        "round": 0,
                    }, ensure_ascii=False).encode('utf-8'))
                    return

                last = json.loads(lines[-1])
                agents = last.get("agents", [])
                elim_meta = {x.get("agent_id"): x for x in elim_from_state if isinstance(x, dict)}

                active = [a for a in agents if a.get("status") == "active"]
                ranked_active = sorted(active, key=lambda x: x.get("score", 0), reverse=True)

                leaderboard = []
                for idx, a in enumerate(ranked_active, start=1):
                    stats = a.get("stats", {})
                    generation = int(a.get("generation", 0) or 0)
                    leaderboard.append({
                        "rank": idx,
                        "status": "active",
                        "agent_id": a.get("agent_id", ""),
                        "generation": generation,
                        "mutation_count": generation,
                        "score": a.get("score", 0),
                        "tasks_completed": stats.get("tasks_completed", 0),
                        "tasks_failed": stats.get("tasks_failed", 0),
                        "consecutive_failures": stats.get("consecutive_failures", 0),
                    })

                eliminated_raw = [a for a in agents if a.get("status") != "active"]
                eliminated_sorted = sorted(
                    eliminated_raw,
                    key=lambda x: (elim_meta.get(x.get("agent_id"), {}).get("round") or 0),
                    reverse=True,
                )
                eliminated_rows = []
                for a in eliminated_sorted:
                    aid = a.get("agent_id", "")
                    stats = a.get("stats", {})
                    generation = int(a.get("generation", 0) or 0)
                    meta = elim_meta.get(aid, {})
                    reason = meta.get("reason", "")
                    eliminated_rows.append({
                        "status": "eliminated",
                        "agent_id": aid,
                        "generation": generation,
                        "mutation_count": generation,
                        "score": a.get("score", 0),
                        "tasks_completed": stats.get("tasks_completed", 0),
                        "tasks_failed": stats.get("tasks_failed", 0),
                        "eliminated_round": meta.get("round"),
                        "elimination_reason": reason,
                    })

                self.wfile.write(json.dumps({
                    "simulation": simulation_dir.name,
                    "round": last.get("round", 0),
                    "leaderboard": leaderboard,
                    "eliminated": eliminated_rows,
                    "environment": env_payload,
                }, ensure_ascii=False).encode('utf-8'))
            except Exception as e:
                self.wfile.write(json.dumps({"error": str(e), "leaderboard": [], "eliminated": []}, ensure_ascii=False).encode('utf-8'))
            return

        if parsed.path == '/api/agent_profile':
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()

            simulation_id = params.get('simulation', [''])[0]
            agent_id = params.get('agent_id', [''])[0]
            simulation_dir = resolve_simulation_dir(simulation_id)
            if not agent_id:
                self.wfile.write(json.dumps({"error": "missing agent_id"}, ensure_ascii=False).encode('utf-8'))
                return

            mem_name = _agent_memory_dir_name(agent_id)
            base = simulation_dir / "agent_memory" / mem_name
            memory_file = base / "memory.jsonl"
            skills_file = base / "skills.json"
            sops_file = base / "sops.json"

            def read_json(path: Path, default):
                if not path.exists():
                    return default
                try:
                    with open(path, 'r', encoding='utf-8') as f:
                        return json.load(f)
                except Exception:
                    return default

            recent_memory = []
            if memory_file.exists():
                try:
                    with open(memory_file, 'r', encoding='utf-8') as f:
                        lines = f.readlines()[-30:]
                    for line in lines:
                        line = line.strip()
                        if line:
                            recent_memory.append(json.loads(line))
                except Exception:
                    recent_memory = []

            response = {
                "simulation": simulation_dir.name,
                "agent_id": agent_id,
                "memory_dir": mem_name,
                "skills": read_json(skills_file, {}),
                "sops": read_json(sops_file, []),
                "recent_memory": recent_memory
            }
            self.wfile.write(json.dumps(response, ensure_ascii=False).encode('utf-8'))
            return

        if parsed.path == '/api/simulations':
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            sims = [p.name for p in list_simulation_dirs()]
            self.wfile.write(json.dumps({"simulations": sims}, ensure_ascii=False).encode('utf-8'))
            return
        
        # API: 获取事件日志（支持增量读取）
        if parsed.path == '/api/events':
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            
            # 获取offset参数
            offset = int(params.get('offset', ['0'])[0])
            simulation_id = params.get('simulation', [''])[0]
            simulation_dir = resolve_simulation_dir(simulation_id)
            events_file = simulation_dir / 'events.jsonl'
            if not events_file.exists():
                self.wfile.write(json.dumps({
                    'events': [],
                    'total': 0,
                    'offset': 0,
                    'simulation': simulation_dir.name
                }, ensure_ascii=False).encode('utf-8'))
                return
            
            # 读取所有行
            with open(events_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 返回offset之后的事件
            new_events = []
            for line in lines[offset:]:
                if line.strip():
                    new_events.append(json.loads(line))
            
            response = {
                'events': new_events,
                'total': len(lines),
                'offset': offset + len(new_events),
                'simulation': simulation_dir.name
            }
            self.wfile.write(json.dumps(response, ensure_ascii=False).encode('utf-8'))
            return
        
        # 默认处理静态文件
        super().do_GET()

if __name__ == '__main__':
    with socketserver.TCPServer(("", PORT), RealTimeHandler) as httpd:
        print(f"🌐 实时监控服务器启动: http://localhost:{PORT}")
        print(f"📊 打开浏览器访问: http://localhost:{PORT}/visualization_realtime.html")
        print("按 Ctrl+C 停止服务器")
        httpd.serve_forever()