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 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


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

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


# ===== 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)


# ===== 可选：配置 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"),
    experiences_path="./.cache/experiences.md",  # 经验文件路径
)
set_runner(runner)


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

# 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)

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