api.py

#!/usr/bin/env python3
"""FastAPI 接口层 — REST API + SSE 流式输出"""

import time
import json
from typing import Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from main import (
    get_full_answer_sync,
    get_session_info,
    get_session_history,
    get_cache_stats,
    get_all_sessions_info,
    clear_session_history,
    generate_response,
    astream_response,
)
from security_module import SecurityManager, validate_session_id

security = SecurityManager()


# ── Lifespan ──────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    from logger import logger
    logger.info("FastAPI 服务启动")
    yield
    logger.info("FastAPI 服务关闭")


app = FastAPI(
    title="茶叶智能客服 API",
    description="基于 LangGraph 的 RAG + Agent 智能客服系统",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Schemas ───────────────────────────────────────────────
class ChatRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=1000, description="用户消息")
    session_id: str = Field("default", max_length=64, description="会话 ID")


class ChatResponse(BaseModel):
    reply: str
    session_id: str
    cached: bool = False
    cache_level: Optional[str] = None
    elapsed_ms: float


# ── Endpoints ─────────────────────────────────────────────

@app.get("/health")
async def health():
    """健康检查"""
    from main import CHUNKS
    return {
        "status": "healthy",
        "chunks": len(CHUNKS),
        "cache_hit_rate": get_cache_stats().get("hit_rate", 0),
    }


@app.post("/chat", response_model=ChatResponse)
async def chat(req: ChatRequest, request: Request):
    """同步聊天接口"""
    if not validate_session_id(req.session_id):
        raise HTTPException(400, "无效的 session_id")

    client_ip = request.client.host if request.client else None
    allowed, reason, _severity = security.check_request(
        req.session_id, ip_address=client_ip, query=req.query
    )
    if not allowed:
        raise HTTPException(429, f"请求被限制: {reason}")

    security.record_request(req.session_id, req.query, client_ip)

    t0 = time.time()
    reply = get_full_answer_sync(req.query, req.session_id)
    elapsed = (time.time() - t0) * 1000

    return ChatResponse(
        reply=reply,
        session_id=req.session_id,
        elapsed_ms=round(elapsed, 1),
    )


@app.post("/chat/stream")
async def chat_stream(req: ChatRequest, request: Request):
    """SSE 流式聊天接口"""
    if not validate_session_id(req.session_id):
        raise HTTPException(400, "无效的 session_id")

    client_ip = request.client.host if request.client else None
    allowed, reason, _severity = security.check_request(
        req.session_id, ip_address=client_ip, query=req.query
    )
    if not allowed:
        raise HTTPException(429, f"请求被限制: {reason}")

    security.record_request(req.session_id, req.query, client_ip)

    async def event_stream():
        try:
            async for token in astream_response(req.query, req.session_id):
                yield f"data: {json.dumps({'token': token}, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
        except Exception as e:
            yield f"data: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/cache/stats")
async def cache_stats():
    """缓存命中统计"""
    return get_cache_stats()


@app.get("/sessions")
async def list_sessions():
    """列出所有活跃会话"""
    return get_all_sessions_info()


@app.get("/sessions/{session_id}")
async def session_info(session_id: str):
    """查询单个会话"""
    if not validate_session_id(session_id):
        raise HTTPException(400, "无效的 session_id")
    info = get_session_info(session_id)
    info["history"] = get_session_history(session_id)
    return info


@app.delete("/sessions/{session_id}")
async def clear_session(session_id: str):
    """清空会话历史"""
    if not validate_session_id(session_id):
        raise HTTPException(400, "无效的 session_id")
    ok = clear_session_history(session_id)
    return {"status": "ok" if ok else "error", "session_id": session_id}


# ── 直接启动 ──────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True)