agent_backend/examples/main.py

"""API 服务开发课程：同一个内存商品服务，练习 HTTP、SSE 与 WebSocket。"""

import asyncio
import json
import os
import secrets
from contextlib import asynccontextmanager
from typing import Annotated

import httpx
from fastapi import Depends, FastAPI, Header, HTTPException, Path, Query, Request
from fastapi import Response, WebSocket, WebSocketDisconnect
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, ConfigDict, Field
from starlette.exceptions import HTTPException as StarletteHTTPException


class ErrorDetail(BaseModel):
    loc: list[str | int]
    msg: str
    type: str


class ErrorContent(BaseModel):
    code: str
    message: str
    details: list[ErrorDetail]


class ErrorResponse(BaseModel):
    error: ErrorContent


@asynccontextmanager
async def lifespan(app: FastAPI):
    timeout = httpx.Timeout(3.0, connect=1.0)
    async with httpx.AsyncClient(
        base_url="http://127.0.0.1:8000", timeout=timeout,
    ) as client:
        app.state.http = client
        yield


app = FastAPI(
    title="API Lab · 商品服务",
    version="1.0.0",
    lifespan=lifespan,
    responses={422: {"model": ErrorResponse, "description": "请求参数校验失败"}},
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5500", "http://127.0.0.1:5500"],
    allow_credentials=False,
    allow_methods=["GET", "POST", "PUT", "DELETE"],
    allow_headers=["Content-Type", "X-API-Key"],
)

# 教学用内存存储：重启后丢失；单进程使用，不能用多个 worker 共享数据。
app.state.products = {}
app.state.next_id = 1


class ProductCreate(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    name: str = Field(min_length=1, max_length=120)
    description: str = Field(default="", max_length=16000)
    is_active: bool = False
    price_cents: int = Field(default=0, ge=0, le=2147483647, strict=True)


class ProductReplace(BaseModel):
    """PUT 完整替换：四个可修改字段都必须提交。"""

    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    name: str = Field(min_length=1, max_length=120)
    description: str = Field(max_length=16000)
    is_active: bool
    price_cents: int = Field(ge=0, le=2147483647, strict=True)


class ProductOut(ProductCreate):
    id: int


class PromptIn(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")

    prompt: str = Field(min_length=1, max_length=200)


@app.exception_handler(StarletteHTTPException)
async def http_error_handler(request: Request, exc: StarletteHTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content={"error": {"code": "http_error", "message": str(exc.detail), "details": []}},
        headers=exc.headers,
    )


@app.exception_handler(RequestValidationError)
async def validation_error_handler(request: Request, exc: RequestValidationError):
    # 只返回定位问题需要的字段，避免回传原始请求、内部异常或不可序列化的 ctx。
    details = [
        {"loc": list(error["loc"]), "msg": error["msg"], "type": error["type"]}
        for error in exc.errors()
    ]
    return JSONResponse(
        status_code=422,
        content={"error": {"code": "validation_error", "message": "请求参数校验失败", "details": details}},
    )


def find_product(product_id: int) -> ProductOut:
    product = app.state.products.get(product_id)
    if product is None:
        raise HTTPException(status_code=404, detail="Product not found")
    return product


@app.get("/health", tags=["起步"])
async def health():
    return {"status": "ok"}


@app.get("/products", response_model=list[ProductOut], tags=["商品 CRUD"])
async def list_products(
    q: Annotated[str | None, Query(max_length=120)] = None,
    limit: Annotated[int, Query(ge=1, le=100)] = 10,
    is_active: bool | None = None,
):
    products = list(app.state.products.values())
    if q:
        query = q.casefold()
        products = [product for product in products if query in product.name.casefold() or query in product.description.casefold()]
    if is_active is not None:
        products = [product for product in products if product.is_active == is_active]
    return products[:limit]


@app.post("/products", response_model=ProductOut, status_code=201, tags=["商品 CRUD"])
async def create_product(body: ProductCreate, response: Response):
    product = ProductOut(id=app.state.next_id, **body.model_dump())
    app.state.products[product.id] = product
    app.state.next_id += 1
    response.headers["Location"] = f"/products/{product.id}"
    return product


@app.get("/products/{product_id}", response_model=ProductOut, responses={404: {"model": ErrorResponse}}, tags=["商品 CRUD"])
async def get_product(product_id: Annotated[int, Path(ge=1)]):
    return find_product(product_id)


@app.put("/products/{product_id}", response_model=ProductOut, responses={404: {"model": ErrorResponse}}, tags=["商品 CRUD"])
async def replace_product(product_id: Annotated[int, Path(ge=1)], body: ProductReplace):
    find_product(product_id)
    product = ProductOut(id=product_id, **body.model_dump())
    app.state.products[product_id] = product
    return product


@app.delete("/products/{product_id}", status_code=204, response_class=Response, responses={404: {"model": ErrorResponse}}, tags=["商品 CRUD"])
async def delete_product(product_id: Annotated[int, Path(ge=1)]):
    find_product(product_id)
    del app.state.products[product_id]
    return Response(status_code=204)


async def require_api_key(x_api_key: Annotated[str | None, Header(alias="X-API-Key")] = None):
    expected = os.getenv("DEMO_API_KEY")
    if not expected:
        raise HTTPException(status_code=503, detail="Set DEMO_API_KEY before using /private")
    if x_api_key is None or not secrets.compare_digest(x_api_key.encode(), expected.encode()):
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing API key",
            headers={"WWW-Authenticate": "APIKey"},
        )


@app.get(
    "/private",
    dependencies=[Depends(require_api_key)],
    responses={401: {"model": ErrorResponse}, 503: {"model": ErrorResponse}},
    tags=["依赖与请求头"],
)
async def private_example():
    return {"message": "API key accepted"}


@app.get("/slow", tags=["异步等待"])
async def slow(delay: Annotated[float, Query(ge=0, le=5)] = 1):
    await asyncio.sleep(delay)
    return {"delay": delay, "status": "done"}


@app.get(
    "/upstream-demo",
    responses={502: {"model": ErrorResponse}, 504: {"model": ErrorResponse}},
    tags=["调用其他 HTTP 接口"],
)
async def upstream_demo(
    request: Request,
    delay: Annotated[float, Query(ge=0, le=5)] = 1,
):
    # 固定调用本机的 /slow，不接受用户提供的出站 URL。
    try:
        response = await request.app.state.http.get("/slow", params={"delay": delay})
        response.raise_for_status()
        return {"upstream": response.json()}
    except httpx.TimeoutException as exc:
        raise HTTPException(504, "上游响应超时") from exc
    except httpx.HTTPStatusError as exc:
        raise HTTPException(502, "上游返回非成功状态") from exc
    except httpx.RequestError as exc:
        raise HTTPException(502, "无法连接上游服务") from exc


@app.post("/generate", tags=["普通响应与流式响应"])
async def generate(payload: PromptIn):
    # 商品清单导出进度的技术演示：模拟三个步骤，不生成实际导出文件。
    parts = []
    for text in ["准备导出商品清单：", payload.prompt, "。导出进度演示完成。"]:
        await asyncio.sleep(0.4)
        parts.append(text)
    return {"text": "".join(parts)}


def encode_sse(event: str, data: dict) -> str:
    payload = json.dumps(data, ensure_ascii=False)
    return f"event: {event}\ndata: {payload}\n\n"


async def token_events(request: Request, prompt: str):
    try:
        for text in ["准备导出商品清单：", prompt, "。导出进度演示完成。"]:
            await asyncio.sleep(0.4)
            if await request.is_disconnected():
                return
            yield encode_sse("token", {"text": text})
        yield encode_sse("done", {})
    finally:
        print("stream closed")


@app.get("/stream", tags=["普通响应与流式响应"])
async def stream(
    request: Request,
    prompt: Annotated[str, Query(min_length=1, max_length=200)] = "catalog",
):
    return StreamingResponse(
        token_events(request, prompt),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.websocket("/ws")
async def websocket_echo(websocket: WebSocket):
    # 商品咨询传输演示：回显收到的文本，不实现客服逻辑或消息存储。
    await websocket.accept()
    try:
        while True:
            message = await websocket.receive_text()
            if message == "ping":
                await websocket.send_json({"type": "pong"})
            else:
                await websocket.send_json({"type": "echo", "text": message})
    except WebSocketDisconnect:
        return