1from fastapi import FastAPI
2from pydantic import BaseModel
3import time
4
5app = FastAPI()
6
7class GenerateRequest(BaseModel):
8    prompt: str
9    max_tokens: int = 128
10
11class GenerateResponse(BaseModel):
12    text: str
13    latency_ms: int
14
15def generate_text(prompt: str, max_tokens: int) -> str:
16    trimmed = " ".join(prompt.split())[:max_tokens]
17    return f"model response for: {trimmed}"
18
19@app.post("/generate", response_model=GenerateResponse)
20def generate(req: GenerateRequest) -> GenerateResponse:
21    started = time.perf_counter()
22    text = generate_text(req.prompt, req.max_tokens)
23    return GenerateResponse(text=text, latency_ms=int((time.perf_counter() - started) * 1000))