Agent/api_server.py

"""
API Server - FastAPI 应用入口

聚合所有模块的 API 路由：
- GET  /api/traces — 查询（trace/api.py）
- POST /api/traces — 执行控制（trace/run_api.py，需配置 Runner）
- WS   /api/traces/{id}/watch — 实时推送（trace/websocket.py）
- GET  /api/experiences — 经验查询（trace/run_api.py，需配置 Runner）
"""

import logging
import json
import os
from dotenv import load_dotenv
load_dotenv()
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.run_api import router as run_router, experiences_router, set_runner
from agent.trace.websocket import router as ws_router, set_trace_store as set_ws_trace_store
from agent.trace.examples_api import router as examples_router
from agent.trace.logs_websocket import router as logs_router, setup_websocket_logging
from agent.trace.upload_api import router as upload_router, set_trace_store as set_upload_trace_store


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

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

# 设置WebSocket日志推送
setup_websocket_logging(level=logging.INFO)


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

app = FastAPI(
    title="Agent API",
    description="Agent 查询 + 执行 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)
set_upload_trace_store(trace_store)


# ===== 可选：配置 Runner（启用执行 API）=====

# 如需启用 POST /api/traces（新建/运行/停止/反思），取消以下注释并配置 LLM：

from agent.core.runner import AgentRunner
from agent.llm import create_openrouter_llm_call

runner = AgentRunner(
    trace_store=trace_store,
    llm_call=create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5"),
)
set_runner(runner)


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

# Examples API（GET /api/examples）
app.include_router(examples_router)

# Trace 上传 API（POST /api/traces/upload）
app.include_router(upload_router)

# Trace 执行 API（POST + GET /running，需配置 Runner）
# 注意：run_router 必须在 api_router 之前注册，否则 GET /running 会被 /{trace_id} 捕获
app.include_router(run_router)

# 经验 API（GET /api/experiences，需配置 Runner）
app.include_router(experiences_router)

# Trace 查询 API（GET）
app.include_router(api_router)

# Trace WebSocket（实时推送）
app.include_router(ws_router)

# Logs WebSocket（日志推送）
app.include_router(logs_router)


@app.on_event("startup")
async def on_startup():
    """服务器启动时执行状态对齐"""
    from agent.trace.run_api import reconcile_traces
    await reconcile_traces()

@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"
    )