Agent/api_server.py

"""
API Server - FastAPI 应用入口

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

import logging
from fastapi import FastAPI
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.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"
    )