Agent/api_server.py

"""
API Server - FastAPI 应用入口

聚合所有模块的 API 路由（step_tree、未来的 memory 等）
"""

import logging
import json
import os
from fastapi import FastAPI, Request, WebSocket
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

from agent.trace import FileSystemTraceStore
from agent.trace.api import router as api_router, set_trace_store as set_api_trace_store
from agent.trace.websocket import router as ws_router, set_trace_store as set_ws_trace_store


# ===== 日志配置 =====

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)


# ===== FastAPI 应用 =====

app = FastAPI(
    title="Agent Step Tree API",
    description="Step 树可视化 API",
    version="1.0.0"
)

# CORS 配置（允许前端跨域访问）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应限制具体域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ===== 初始化存储 =====

# 使用文件系统存储（支持跨进程和持久化）
trace_store = FileSystemTraceStore(base_path=".trace")

# 注入到 step_tree 模块
set_api_trace_store(trace_store)
set_ws_trace_store(trace_store)


# ===== 注册路由 =====

# Step 树 RESTful API
app.include_router(api_router)

# Step 树 WebSocket
app.include_router(ws_router)

@app.websocket("/ws_ping")
async def ws_ping(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_text("pong")
    await websocket.close()


# ===== 健康检查 =====

@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "ok",
        "service": "Agent Step Tree API",
        "version": "1.0.0"
    }


# ===== 启动服务 =====

if __name__ == "__main__":
    logger.info("Starting API server...")
    uvicorn.run(
        "api_server:app",
        host="0.0.0.0",
        port=8000,
        reload=True,  # 开发模式
        log_level="info"
    )