Real-time AI Systems

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from openai import OpenAI
import json, time

app = FastAPI()
client = OpenAI()

@app.post("/chat/stream")
async def stream_chat(prompt: str):
    """Stream LLM tokens via Server-Sent Events."""
    async def generate():
        start = time.time()
        stream = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            stream=True,
            stream_options={"include_usage": True},
        )
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                # SSE format: "data: <json>\n\n"
                payload = {
                    "token": chunk.choices[0].delta.content,
                    "latency_ms": round((time.time() - start) * 1000),
                }
                yield f"data: {json.dumps(payload)}\n\n"

            # Final chunk includes usage stats
            if chunk.usage:
                yield f"data: {json.dumps({"usage": chunk.usage.model_dump()})}\n\n"

        yield "data: [DONE]\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",  # Disable nginx buffering
        },
    )

// Browser: consuming SSE with fetch + ReadableStream
async function streamChat(prompt: string, onToken: (t: string) => void) {
  const response = await fetch("/chat/stream", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ prompt }),
  });

  const reader = response.body!.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split("\n\n");
    buffer = lines.pop()!; // Keep incomplete chunk

    for (const line of lines) {
      const data = line.replace("data: ", "");
      if (data === "[DONE]") return;
      const { token } = JSON.parse(data);
      if (token) onToken(token); // Append to UI
    }
  }
}

// Usage: tokens appear character by character
streamChat("Explain transformers", (token) => {
  document.getElementById("output")!.textContent += token;
});

# Start vLLM with production settings
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \           # Shard across 4 GPUs
    --max-model-len 8192 \               # Max context length
    --enable-prefix-caching \            # Reuse KV for shared prefixes
    --quantization awq \                 # 4-bit quantization
    --gpu-memory-utilization 0.9 \       # Use 90% of GPU memory
    --max-num-seqs 256 \                 # Max concurrent sequences
    --port 8000

# Then use with any OpenAI-compatible client
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="unused")

# Streaming response
stream = client.chat.completions.create(
    model="meta-llama/Llama-3.1-70B-Instruct",
    messages=[{"role": "user", "content": "Explain PagedAttention"}],
    stream=True,
    max_tokens=512,
)
for chunk in stream:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)