S130 · Design an AI Chat Product — RAG + Agents + Serving

"""
Minimal AI chat orchestrator — one turn, honest streaming.
Deps: fastapi uvicorn openai qdrant-client tiktoken
Run:  uvicorn s130_chat:app --reload --port 8000
"""
from __future__ import annotations
 
import asyncio
import json
import os
import time
from typing import AsyncIterator
 
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI
from qdrant_client import AsyncQdrantClient
 
app = FastAPI()
openai = AsyncOpenAI(api_key=os.environ["OPENAI_API_KEY"])
vdb = AsyncQdrantClient(url=os.environ.get("QDRANT_URL", "http://localhost:6333"))
 
FLAGSHIP = "gpt-4o"
CHEAP = "gpt-4o-mini"
DAILY_TOKEN_CAP = 200_000  # per user
SYSTEM_PROMPT = "You are a helpful assistant. Cite sources when the answer comes from retrieved context."
 
# --- fake in-memory user state; replace with Redis in prod --------------------
user_tokens_today: dict[str, int] = {}
conversations: dict[str, list[dict]] = {}
 
 
async def route(user_msg: str) -> str:
    """Cheap model classifies whether we need RAG. Returns 'rag' or 'direct'."""
    r = await openai.chat.completions.create(
        model=CHEAP,
        messages=[
            {"role": "system", "content": "Reply with exactly one word: 'rag' if the user's question likely needs external facts about a specific company, product, doc or event; 'direct' otherwise."},
            {"role": "user", "content": user_msg},
        ],
        max_tokens=2,
        temperature=0,
    )
    return r.choices[0].message.content.strip().lower()
 
 
async def retrieve(user_msg: str, k: int = 3) -> list[str]:
    """Embed + vector-DB top-k. Skip re-ranker here for brevity."""
    emb = (await openai.embeddings.create(model="text-embedding-3-small", input=user_msg)).data[0].embedding
    hits = await vdb.search(collection_name="docs", query_vector=emb, limit=k)
    return [h.payload["text"] for h in hits]
 
 
def enforce_budget(user_id: str, projected: int) -> None:
    used = user_tokens_today.get(user_id, 0)
    if used + projected > DAILY_TOKEN_CAP:
        raise HTTPException(status_code=429, detail="Daily token budget exceeded")
 
 
async def stream_turn(user_id: str, user_msg: str) -> AsyncIterator[str]:
    enforce_budget(user_id, projected=1_000)  # rough upper bound per turn
 
    route_decision = await route(user_msg)
    yield sse({"event": "route", "value": route_decision})
 
    context_blocks: list[str] = []
    if route_decision == "rag":
        context_blocks = await retrieve(user_msg)
        yield sse({"event": "sources", "count": len(context_blocks)})
 
    history = conversations.setdefault(user_id, [])
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    if context_blocks:
        messages.append({"role": "system", "content": "Context:\n\n" + "\n---\n".join(context_blocks)})
    messages.extend(history[-20:])  # cap history
    messages.append({"role": "user", "content": user_msg})
 
    reply_parts: list[str] = []
    completion_tokens = 0
    async with openai.chat.completions.stream(
        model=FLAGSHIP, messages=messages, max_tokens=800
    ) as stream:
        async for event in stream:
            if event.type == "content.delta":
                reply_parts.append(event.delta)
                completion_tokens += 1
                yield sse({"event": "token", "value": event.delta})
 
    reply = "".join(reply_parts)
    history.append({"role": "user", "content": user_msg})
    history.append({"role": "assistant", "content": reply})
    user_tokens_today[user_id] = user_tokens_today.get(user_id, 0) + completion_tokens
 
    # kick off async hallucination check — never blocks the response
    asyncio.create_task(post_turn_eval(user_id, user_msg, reply, context_blocks))
 
    yield sse({"event": "done", "tokens": completion_tokens})
 
 
async def post_turn_eval(user_id: str, q: str, a: str, ctx: list[str]) -> None:
    """LLM-as-judge hallucination check. Runs after the response ships."""
    if not ctx:
        return
    verdict = await openai.chat.completions.create(
        model=CHEAP,
        messages=[
            {"role": "system", "content": "Reply with 'ok' if the assistant answer is supported by the context; 'suspect' otherwise."},
            {"role": "user", "content": f"Context:\n{ctx}\n\nQ: {q}\nA: {a}"},
        ],
        max_tokens=2,
        temperature=0,
    )
    tag = verdict.choices[0].message.content.strip().lower()
    # emit to your observability sink — LangSmith / Phoenix / Datadog
    print(f"[eval] user={user_id} verdict={tag}")
 
 
def sse(payload: dict) -> str:
    return f"data: {json.dumps(payload)}\n\n"
 
 
@app.post("/chat")
async def chat(request: Request):
    body = await request.json()
    return StreamingResponse(
        stream_turn(body["user_id"], body["message"]),
        media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no", "Cache-Control": "no-cache"},
    )