Pythonmain.py
Copy
1
from fastapi import FastAPI, HTTPException
2
from fastapi.responses import StreamingResponse
3
from pydantic import BaseModel
4
from openai import AsyncOpenAI
5
import asyncio
6
 
7
app = FastAPI(title="LLM Serving API")
8
client = AsyncOpenAI()
9
 
10
class Message(BaseModel):
11
    role: str
12
    content: str
13
 
14
class ChatRequest(BaseModel):
15
    model: str = "gpt-4o-mini"
16
    messages: list[Message]
17
    stream: bool = False
18
 
19
@app.get("/health")
20
async def health():
21
    return {"status": "ok"}
22
 
23
@app.post("/chat")
24
async def chat(req: ChatRequest):
25
    messages = [m.model_dump() for m in req.messages]
26
 
27
    if not req.stream:
28
        resp = await client.chat.completions.create(
29
            model=req.model, messages=messages
30
        )
31
        return {"content": resp.choices[0].message.content}
32
 
33
    async def event_generator():
34
        async with await client.chat.completions.create(
35
            model=req.model, messages=messages, stream=True
36
        ) as stream:
37
            async for chunk in stream:
38
                delta = chunk.choices[0].delta.content
39
                if delta:
40
                    yield f"data: {delta}\n\n"
41
        yield "data: [DONE]\n\n"
42
 
43
    return StreamingResponse(event_generator(), media_type="text/event-stream")
44