"""FastAPI 模拟流式生成：演示超时、并发门控与资源释放。"""
import asyncio
from collections.abc import AsyncIterator
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

app = FastAPI()
slots = asyncio.Semaphore(4)


class Request(BaseModel):
    prompt: str = Field(min_length=1, max_length=2_000)
    max_new_tokens: int = Field(default=32, ge=1, le=256)


async def fake_generate(req: Request) -> AsyncIterator[str]:
    async with slots:
        for index in range(req.max_new_tokens):
            await asyncio.sleep(0.02)
            yield f"data: token-{index}\n\n"
        yield "event: done\ndata: {}\n\n"


@app.post("/generate")
async def generate(req: Request):
    try:
        await asyncio.wait_for(slots.acquire(), timeout=2.0)
        slots.release()  # 这里只做排队探测；真正占用发生在生成器内部。
    except TimeoutError as exc:
        raise HTTPException(status_code=503, detail="server busy") from exc
    return StreamingResponse(fake_generate(req), media_type="text/event-stream")