noke.py

#!/usr/bin/env python3

import os
import logging
import asyncio
from pathlib import Path
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, RedirectResponse

from core.config import ConfigManager
from core.token_manager import TokenManager
from core.log_store import LogStore
from core.conversation_cache import ConversationCache
from core.ratelimit import RateLimiter
from routes import openai_compat, admin_api, claude_api, system_api, mcp_api

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

cfg = ConfigManager()
token_manager = TokenManager(cfg)
log_store = LogStore(max_entries=cfg.get("logging", "max_entries", default=500))

# 会话缓存 & 限流器
conv_cache = ConversationCache(max_size=2000, ttl=1800)
limiter = RateLimiter(
    qps=cfg.get("proxy", "rate_limit_qps", default=0),
    burst=cfg.get("proxy", "rate_limit_burst", default=0),
)

# 初始化各路由模块
openai_compat.init(token_manager, cfg, log_store, conv_cache, limiter)
admin_api.init(cfg, token_manager, log_store)
claude_api.init(token_manager, cfg, log_store, conv_cache, limiter)
system_api.init(token_manager, cfg, conv_cache, limiter)
mcp_api.init(token_manager, cfg)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(
        "Noke started — tokens: %d, port: %d",
        len(cfg.get("tokens", default=[])),
        int(os.environ.get("PORT", cfg.get("server", "port", default=7860))),
    )

    # 后台定时 flush token 统计（每 60 秒，标脏才写盘）
    async def _periodic_flush():
        while True:
            await asyncio.sleep(60)
            token_manager.flush()
    flush_task = asyncio.create_task(_periodic_flush())

    yield

    flush_task.cancel()
    try:
        await flush_task
    except asyncio.CancelledError:
        pass
    await token_manager.close_all()


app = FastAPI(lifespan=lifespan)

# 全局异常处理：脱敏
@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    import traceback
    error_detail = f"{type(exc).__name__}: {str(exc)}\n{traceback.format_exc()}"
    logger.error("Unhandled exception: %s", error_detail)
    from fastapi.responses import JSONResponse
    return JSONResponse(
        status_code=500,
        content={"detail": "internal server error"}
    )

app.include_router(claude_api.router)
app.include_router(openai_compat.router)
app.include_router(admin_api.router)
app.include_router(system_api.router)
app.include_router(mcp_api.router)

static_dir = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")


@app.get("/")
async def root():
    return RedirectResponse(url="/admin")


# health 现在由 system_api 提供，保留旧路径兼容
@app.get("/health")
async def health_legacy():
    return {"status": "ok"}


@app.get("/admin")
async def admin_page():
    return FileResponse(str(static_dir / "index.html"))


if __name__ == "__main__":
    import urllib3
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(
        app,
        host=cfg.get("server", "host", default="0.0.0.0"),
        port=port,
    )