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.notes = {}
app.state.next_id = 1


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

    title: str = Field(min_length=1, max_length=120)
    content: str = ""
    pinned: bool = False


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

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

    title: str = Field(min_length=1, max_length=120)
    content: str
    pinned: bool


class NoteOut(NoteCreate):
    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_note(note_id: int) -> NoteOut:
    note = app.state.notes.get(note_id)
    if note is None:
        raise HTTPException(status_code=404, detail="Note not found")
    return note


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


@app.get("/notes", response_model=list[NoteOut], tags=["笔记 CRUD"])
async def list_notes(
    q: Annotated[str | None, Query(max_length=120)] = None,
    limit: Annotated[int, Query(ge=1, le=100)] = 10,
    pinned: bool | None = None,
):
    notes = list(app.state.notes.values())
    if q:
        query = q.casefold()
        notes = [note for note in notes if query in note.title.casefold() or query in note.content.casefold()]
    if pinned is not None:
        notes = [note for note in notes if note.pinned == pinned]
    return notes[:limit]


@app.post("/notes", response_model=NoteOut, status_code=201, tags=["笔记 CRUD"])
async def create_note(body: NoteCreate, response: Response):
    note = NoteOut(id=app.state.next_id, **body.model_dump())
    app.state.notes[note.id] = note
    app.state.next_id += 1
    response.headers["Location"] = f"/notes/{note.id}"
    return note


@app.get("/notes/{note_id}", response_model=NoteOut, responses={404: {"model": ErrorResponse}}, tags=["笔记 CRUD"])
async def get_note(note_id: Annotated[int, Path(ge=1)]):
    return find_note(note_id)


@app.put("/notes/{note_id}", response_model=NoteOut, responses={404: {"model": ErrorResponse}}, tags=["笔记 CRUD"])
async def replace_note(note_id: Annotated[int, Path(ge=1)], body: NoteReplace):
    find_note(note_id)
    note = NoteOut(id=note_id, **body.model_dump())
    app.state.notes[note_id] = note
    return note


@app.delete("/notes/{note_id}", status_code=204, response_class=Response, responses={404: {"model": ErrorResponse}}, tags=["笔记 CRUD"])
async def delete_note(note_id: Annotated[int, Path(ge=1)]):
    find_note(note_id)
    del app.state.notes[note_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):
    # 确定性输出，无模型 SDK 或模型密钥；等待三步完成后一次返回。
    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)] = "hello",
):
    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