FastAPI 快速上手（AI 服务接口）¶

from __future__ import annotations

import asyncio
import json
import logging
from collections.abc import AsyncIterator
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

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ai-service")


class PredictRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=10_000)
    model: str = Field(default="demo-llm")
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=256, ge=1, le=4096)


class PredictResponse(BaseModel):
    text: str
    usage: dict[str, int]
    model: str


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatStreamRequest(BaseModel):
    messages: list[ChatMessage]


class MockAIModel:
    """教学用 mock；真实项目里换成 vLLM / transformers / ONNX Runtime 等实现。"""

    async def generate(
        self,
        prompt: str,
        *,
        temperature: float,
        max_tokens: int,
    ) -> dict[str, object]:
        await asyncio.sleep(0.05)
        output = f"AI 回复：{prompt[:80]}"
        return {
            "text": output,
            "usage": {
                "prompt_tokens": min(len(prompt.split()), max_tokens),
                "completion_tokens": min(32, max_tokens),
            },
        }

    async def stream_text(self, prompt: str) -> AsyncIterator[str]:
        for token in prompt.split():
            await asyncio.sleep(0.03)
            yield token


def load_model() -> MockAIModel:
    return MockAIModel()


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.model = load_model()
    logger.info("Model loaded")
    yield
    del app.state.model
    logger.info("Model released")


app = FastAPI(title="AI Service API", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://your-frontend.example.com"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

inference_semaphore = asyncio.Semaphore(4)


@app.middleware("http")
async def log_requests(request: Request, call_next):
    start = asyncio.get_running_loop().time()
    response = await call_next(request)
    duration = asyncio.get_running_loop().time() - start
    logger.info(
        "%s %s -> %s (%.3fs)",
        request.method,
        request.url.path,
        response.status_code,
        duration,
    )
    return response


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/predict", response_model=PredictResponse)
async def predict(req: PredictRequest, request: Request) -> PredictResponse:
    model: MockAIModel = request.app.state.model

    async with inference_semaphore:
        result = await model.generate(
            req.text,
            temperature=req.temperature,
            max_tokens=req.max_tokens,
        )

    return PredictResponse(
        text=result["text"],
        usage=result["usage"],
        model=req.model,
    )


async def sse_events(model: MockAIModel, prompt: str) -> AsyncIterator[str]:
    async for token in model.stream_text(prompt):
        payload = json.dumps({"token": token}, ensure_ascii=False)
        yield f"data: {payload}\n\n"
    yield "event: done\ndata: [DONE]\n\n"


@app.post("/chat/stream")
async def chat_stream(req: ChatStreamRequest, request: Request) -> StreamingResponse:
    if not req.messages:
        raise HTTPException(status_code=400, detail="messages 不能为空")

    prompt = "\n".join(item.content for item in req.messages)
    model: MockAIModel = request.app.state.model
    return StreamingResponse(sse_events(model, prompt), media_type="text/event-stream")

async def sse_events(model: MockAIModel, prompt: str) -> AsyncIterator[str]:
    async for token in model.stream_text(prompt):
        yield f"data: {json.dumps({'token': token}, ensure_ascii=False)}\n\n"

    yield "event: done\ndata: [DONE]\n\n"

import io

from fastapi import File, UploadFile
from PIL import Image


def run_ocr(image: Image.Image) -> str:
    # 这里换成真实 OCR 模型调用
    return f"image-size={image.size}"


@app.post("/ocr")
async def ocr(file: UploadFile = File(...)) -> dict[str, str]:
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    return {"text": run_ocr(image)}

inference_semaphore = asyncio.Semaphore(4)


@app.post("/predict", response_model=PredictResponse)
async def predict(req: PredictRequest, request: Request) -> PredictResponse:
    model: MockAIModel = request.app.state.model

    async with inference_semaphore:
        result = await model.generate(
            req.text,
            temperature=req.temperature,
            max_tokens=req.max_tokens,
        )

    return PredictResponse(text=result["text"], usage=result["usage"], model=req.model)