Streaming LLM Responses

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio, json

app = FastAPI()

async def generate_stream(prompt: str):
    # Simulate LLM streaming (replace with real model)
    for token in ["Hello", " there", "!", " How", " can", " I", " help", "?"]:
        data = json.dumps({"token": token})
        yield f"data: {data}\n\n"
        await asyncio.sleep(0.1)
    yield "data: [DONE]\n\n"

@app.post("/chat/stream")
async def stream_chat(prompt: str):
    return StreamingResponse(
        generate_stream(prompt),
        media_type="text/event-stream"
    )

from openai import AsyncOpenAI

client = AsyncOpenAI()

async def openai_stream(prompt: str):
    stream = await client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )
    async for chunk in stream:
        if chunk.choices[0].delta.content:
            data = json.dumps({"token": chunk.choices[0].delta.content})
            yield f"data: {data}\n\n"
    yield "data: [DONE]\n\n"

@app.post("/chat")
async def chat(prompt: str):
    return StreamingResponse(openai_stream(prompt), media_type="text/event-stream")

from fastapi import WebSocket, WebSocketDisconnect

@app.websocket("/ws/chat")
async def websocket_chat(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            prompt = await ws.receive_text()
            # Stream response back
            async for token in generate_tokens(prompt):
                await ws.send_json({"token": token})
            await ws.send_json({"done": True})
    except WebSocketDisconnect:
        print("Client disconnected")

const response = await fetch("/chat/stream", {
  method: "POST",
  headers: { "Content-Type": "application/json" },
  body: JSON.stringify({ prompt: "Hello" })
});

const reader = response.body.getReader();
const decoder = new TextDecoder();

while (true) {
  const { done, value } = await reader.read();
  if (done) break;
  const text = decoder.decode(value);
  // Parse SSE data and update UI
}

from fastapi import BackgroundTasks

async def log_inference(prompt: str, result: str):
    # Save to database, send metrics, etc.
    await db.save_log(prompt=prompt, result=result)

@app.post("/predict")
async def predict(input: TextInput, background_tasks: BackgroundTasks):
    result = model.predict(input.text)
    background_tasks.add_task(log_inference, input.text, result)
    return {"result": result}  # Returns immediately