স্ট্রিমিংসহ LLM ব্যাকএন্ড বানানো

client --> routes: auth, schema check, daily-limit check   (refuse BEFORE spending)
          --> ChatService: redact, add system prompt, call provider, log, record
                --> RetryingProvider --> PracticeProvider | OpenAIProvider | AnthropicProvider
                --> UsageStore: tokens, cost, status per call

mkdir -p ~/projects/llm-gateway
cd ~/projects/llm-gateway
uv init --bare -q
uv add -q "fastapi[standard]" pydantic-settings httpx2
uv add -q --dev pytest
# with pip: python -m venv .venv && source .venv/bin/activate && pip install -q "fastapi[standard]" pydantic-settings httpx2 pytest
mkdir -p app tests
touch app/__init__.py
cat >> pyproject.toml <<'EOF'

[tool.fastapi]
entrypoint = "app.main:app"
EOF

# file: app/settings.py
from pydantic import BaseModel, SecretStr
from pydantic_settings import BaseSettings, SettingsConfigDict


class ModelConfig(BaseModel):
    name: str = "practice-llm-1"
    max_output_tokens: int = 300
    timeout_s: float = 10.0


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_prefix="LLM_", env_nested_delimiter="__", env_file=".env", extra="ignore")

    provider_url: str = "http://localhost:8800"
    provider_api_key: SecretStr                      # required: no default
    llm: ModelConfig = ModelConfig()
    prompt_version: str = "tutor-v2"
    daily_token_limit: int = 120                     # per user, per UTC day
    price_in_per_m: float = 3.0                      # made-up USD per million tokens
    price_out_per_m: float = 15.0
    api_keys: dict[str, str] = {"gw-nadia": "nadia", "gw-rafi": "rafi", "gw-admin": "admin"}

# file: app/schemas.py
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field


class ChatMessageIn(BaseModel):
    role: Literal["user", "assistant"]            # clients can never send a "system" message
    content: str = Field(min_length=1, max_length=2000)


class ChatIn(BaseModel):
    model_config = ConfigDict(extra="forbid")     # no "model" field: the server decides which model to pay for

    messages: list[ChatMessageIn] = Field(min_length=1, max_length=20)


class ChatOut(BaseModel):
    answer: str
    model: str
    prompt_version: str
    usage: dict[str, int | float]

# file: app/providers.py
import asyncio
import json
import random
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import asdict, dataclass
from typing import Protocol

import httpx

from app.settings import ModelConfig


@dataclass
class Message:
    role: str                    # "system", "user" or "assistant"
    content: str


@dataclass
class Usage:
    prompt_tokens: int
    completion_tokens: int


@dataclass
class Completion:
    text: str
    usage: Usage


@dataclass
class StreamChunk:
    text: str = ""               # a piece of the answer ...
    usage: Usage | None = None   # ... or, last of all, the token counts


class ProviderError(Exception):
    def __init__(self, message: str, *, retryable: bool = False, retry_after: float | None = None,
                 timeout: bool = False) -> None:
        super().__init__(message)
        self.retryable, self.retry_after, self.timeout = retryable, retry_after, timeout


class LLMProvider(Protocol):
    async def complete(self, messages: list[Message], config: ModelConfig) -> Completion: ...
    def stream(self, messages: list[Message], config: ModelConfig) -> AsyncIterator[StreamChunk]: ...
    async def ping(self) -> bool: ...


def sdk_error(exc: Exception, retryable: set[int]) -> ProviderError:
    status = getattr(exc, "status_code", None)            # None: the network failed, not the API
    return ProviderError(f"status {status}", retryable=status is None or status in retryable,
                         timeout=type(exc).__name__ == "APITimeoutError")


def http_error(exc: httpx.HTTPError) -> ProviderError:
    return ProviderError("provider failed", retryable=True, timeout=isinstance(exc, httpx.TimeoutException))


def check_status(response: httpx.Response) -> None:
    if response.status_code >= 400:
        after = response.headers.get("Retry-After", "")
        raise ProviderError(f"status {response.status_code}", retryable=response.status_code in (429, 500, 502, 503, 504),
                            retry_after=float(after) if after.isdigit() else None)


class PracticeProvider:
    def __init__(self, client: httpx.AsyncClient) -> None:
        self.client = client

    @staticmethod
    def _body(messages: list[Message], config: ModelConfig, stream: bool) -> dict:
        return {"model": config.name, "messages": [asdict(m) for m in messages], "stream": stream}

    async def complete(self, messages: list[Message], config: ModelConfig) -> Completion:
        try:
            response = await self.client.post("/v1/chat", json=self._body(messages, config, False))
        except httpx.TransportError as exc:
            raise http_error(exc) from exc
        check_status(response)
        data = response.json()
        return Completion(data["message"]["content"], Usage(data["usage"]["prompt_tokens"], data["usage"]["completion_tokens"]))

    async def stream(self, messages: list[Message], config: ModelConfig) -> AsyncIterator[StreamChunk]:
        try:
            async with self.client.stream("POST", "/v1/chat", json=self._body(messages, config, True)) as response:
                if response.status_code >= 400:
                    await response.aread()
                    check_status(response)
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        data = json.loads(line.removeprefix("data: "))
                        if data.get("done"):
                            yield StreamChunk(usage=Usage(data["usage"]["prompt_tokens"], data["usage"]["completion_tokens"]))
                        else:
                            yield StreamChunk(text=data["delta"])
        except httpx.TransportError as exc:
            raise http_error(exc) from exc

    async def ping(self) -> bool:
        try:
            return (await self.client.get("/health", timeout=2.0)).status_code == 200
        except httpx.HTTPError:
            return False


class RetryingProvider:
    def __init__(self, inner: LLMProvider, *, max_attempts: int = 3, base_delay: float = 0.5,
                 sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
                 jitter: Callable[[], float] = lambda: random.uniform(0.5, 1.0)) -> None:
        self.inner, self.max_attempts, self.base_delay = inner, max_attempts, base_delay
        self.sleep, self.jitter = sleep, jitter

    async def _pause(self, attempt: int, exc: ProviderError) -> None:
        wanted = exc.retry_after if exc.retry_after is not None else self.base_delay * 2 ** (attempt - 1)
        await self.sleep(min(8.0, wanted) * self.jitter())

    async def complete(self, messages: list[Message], config: ModelConfig) -> Completion:
        for attempt in range(1, self.max_attempts + 1):
            try:
                return await self.inner.complete(messages, config)
            except ProviderError as exc:
                if not exc.retryable or attempt == self.max_attempts:
                    raise
                await self._pause(attempt, exc)
        raise AssertionError("unreachable")

    async def stream(self, messages: list[Message], config: ModelConfig) -> AsyncIterator[StreamChunk]:
        for attempt in range(1, self.max_attempts + 1):
            started = False
            try:
                async for chunk in self.inner.stream(messages, config):
                    started = True
                    yield chunk
                return
            except ProviderError as exc:
                if started or not exc.retryable or attempt == self.max_attempts:
                    raise                    # the client already has part of the answer: never start over
                await self._pause(attempt, exc)

    async def ping(self) -> bool:
        return await self.inner.ping()

# app/provider_openai.py   (uv add openai)
from openai import APIConnectionError, APIStatusError, AsyncOpenAI

from app.providers import Completion, ProviderError, StreamChunk, Usage, sdk_error

RETRYABLE = {429, 500, 502, 503, 504}


class OpenAIProvider:
    def __init__(self, api_key: str, timeout_s: float) -> None:
        self.client = AsyncOpenAI(api_key=api_key, timeout=timeout_s, max_retries=0)

    @staticmethod
    def _args(messages, config) -> dict:
        return {
            "model": config.name,
            "instructions": "\n".join(m.content for m in messages if m.role == "system"),
            "input": [{"role": m.role, "content": m.content} for m in messages if m.role != "system"],
            "max_output_tokens": config.max_output_tokens,
        }

    async def complete(self, messages, config) -> Completion:
        try:
            response = await self.client.responses.create(**self._args(messages, config))
        except (APIStatusError, APIConnectionError) as exc:
            raise sdk_error(exc, RETRYABLE) from exc
        return Completion(response.output_text, Usage(response.usage.input_tokens, response.usage.output_tokens))

    async def stream(self, messages, config):
        try:
            events = await self.client.responses.create(**self._args(messages, config), stream=True)
            async for event in events:
                if event.type == "response.output_text.delta":
                    yield StreamChunk(text=event.delta)
                elif event.type == "response.completed":
                    usage = event.response.usage
                    yield StreamChunk(usage=Usage(usage.input_tokens, usage.output_tokens))
                elif event.type == "error":
                    raise ProviderError(event.message)
        except (APIStatusError, APIConnectionError) as exc:
            raise sdk_error(exc, RETRYABLE) from exc

# app/provider_anthropic.py   (uv add anthropic)
from anthropic import APIConnectionError, APIStatusError, AsyncAnthropic

from app.providers import StreamChunk, Usage, sdk_error

RETRYABLE = {429, 500, 502, 503, 504, 529}      # 529: Anthropic is overloaded


class AnthropicProvider:
    def __init__(self, api_key: str, timeout_s: float) -> None:
        self.client = AsyncAnthropic(api_key=api_key, timeout=timeout_s, max_retries=0)

    @staticmethod
    def _args(messages, config) -> dict:
        return {
            "model": config.name,
            "max_tokens": config.max_output_tokens,
            "system": "\n".join(m.content for m in messages if m.role == "system"),
            "messages": [{"role": m.role, "content": m.content} for m in messages if m.role != "system"],
        }

    async def stream(self, messages, config):
        try:
            async with self.client.messages.stream(**self._args(messages, config)) as stream:
                async for text in stream.text_stream:
                    yield StreamChunk(text=text)
                final = await stream.get_final_message()
        except (APIStatusError, APIConnectionError) as exc:
            raise sdk_error(exc, RETRYABLE) from exc
        yield StreamChunk(usage=Usage(final.usage.input_tokens, final.usage.output_tokens))

# file: app/redact.py
import re

PATTERNS = {
    "EMAIL": re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+"),
    "KEY": re.compile(r"\b(?:sk|gw|token)-[A-Za-z0-9_-]{8,}"),
    "PHONE": re.compile(r"(?<!\w)(?:\+?880|0)1\d{9}\b"),        # Bangladesh mobile numbers
}


def redact(text: str) -> str:
    for label, pattern in PATTERNS.items():
        text = pattern.sub(f"[{label}]", text)
    return text

# file: app/service.py
import asyncio
import json
import logging
from collections.abc import AsyncIterator
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime

from app.providers import LLMProvider, Message, ProviderError, Usage
from app.redact import redact
from app.schemas import ChatOut
from app.settings import Settings

logging.basicConfig(format="%(message)s")          # root stays at WARNING, so other libraries stay quiet
log = logging.getLogger("llm")
log.setLevel(logging.INFO)

PROMPTS = {
    "tutor-v1": "You are a BitByte tutor. Answer in one short paragraph.",
    "tutor-v2": "You are a BitByte tutor. Answer in one short paragraph and never repeat personal data.",
}


class GatewayError(Exception):
    def __init__(self, status: int, code: str, message: str, headers: dict[str, str] | None = None) -> None:
        super().__init__(message)
        self.status, self.code, self.message, self.headers = status, code, message, headers or {}


def to_gateway_error(exc: ProviderError | TimeoutError) -> GatewayError:
    if isinstance(exc, TimeoutError) or exc.timeout:
        return GatewayError(504, "provider_timeout", "The model took too long to answer")
    if exc.retryable:
        return GatewayError(503, "provider_busy", "The model provider is busy, try again shortly")
    return GatewayError(502, "provider_error", "The model provider could not handle this request")


@dataclass
class CallRecord:
    user: str
    model: str
    prompt_version: str
    prompt_tokens: int
    completion_tokens: int
    cost_usd: float
    status: str                  # "ok", "error" or "cancelled"
    day: str = field(default_factory=lambda: datetime.now(UTC).date().isoformat())


def totals(records: list[CallRecord]) -> dict:
    return {"requests": len(records), "cancelled": sum(r.status == "cancelled" for r in records),
            "prompt_tokens": sum(r.prompt_tokens for r in records),
            "completion_tokens": sum(r.completion_tokens for r in records),
            "cost_usd": round(sum(r.cost_usd for r in records), 6)}


class UsageStore:
    def __init__(self) -> None:
        self.records: list[CallRecord] = []

    def today(self, user: str) -> list[CallRecord]:
        day = datetime.now(UTC).date().isoformat()
        return [r for r in self.records if r.user == user and r.day == day]

    def tokens_today(self, user: str) -> int:
        return sum(r.prompt_tokens + r.completion_tokens for r in self.today(user))

    def by_prompt_version(self) -> dict[str, dict]:
        groups: dict[str, list[CallRecord]] = {}
        for record in self.records:
            groups.setdefault(f"{record.prompt_version} / {record.model}", []).append(record)
        return {key: totals(records) for key, records in groups.items()}


@dataclass
class Prepared:
    user: str
    messages: list[Message]


class ChatService:
    def __init__(self, provider: LLMProvider, store: UsageStore, settings: Settings) -> None:
        self.provider, self.store, self.settings = provider, store, settings

    def prepare(self, user: str, incoming: list[Message]) -> Prepared:
        if self.store.tokens_today(user) >= self.settings.daily_token_limit:
            now = datetime.now(UTC)
            wait = 86400 - (now.hour * 3600 + now.minute * 60 + now.second)         # seconds until the next UTC day
            raise GatewayError(429, "usage_limit_exceeded", f"Daily limit of {self.settings.daily_token_limit} tokens reached",
                               {"Retry-After": str(wait)})
        system = Message("system", PROMPTS[self.settings.prompt_version])
        return Prepared(user, [system, *(Message(m.role, redact(m.content)) for m in incoming)])

    def _usage_out(self, usage: Usage) -> dict:
        cost = (usage.prompt_tokens * self.settings.price_in_per_m
                + usage.completion_tokens * self.settings.price_out_per_m) / 1_000_000
        return {"prompt_tokens": usage.prompt_tokens, "completion_tokens": usage.completion_tokens,
                "total_tokens": usage.prompt_tokens + usage.completion_tokens, "cost_usd": round(cost, 6)}

    def _record(self, prepared: Prepared, usage: Usage, status: str) -> None:
        out = self._usage_out(usage)
        record = CallRecord(prepared.user, self.settings.llm.name, self.settings.prompt_version, usage.prompt_tokens,
                            usage.completion_tokens, out["cost_usd"], status)
        self.store.records.append(record)
        log.info(json.dumps({"event": "llm_call", **asdict(record)}))

    async def complete(self, prepared: Prepared) -> ChatOut:
        try:
            async with asyncio.timeout(self.settings.llm.timeout_s):
                result = await self.provider.complete(prepared.messages, self.settings.llm)
        except (ProviderError, TimeoutError) as exc:
            self._record(prepared, Usage(0, 0), "error")
            raise to_gateway_error(exc) from exc
        self._record(prepared, result.usage, "ok")
        return ChatOut(answer=result.text, model=self.settings.llm.name, prompt_version=self.settings.prompt_version,
                       usage=self._usage_out(result.usage))

    async def stream(self, prepared: Prepared) -> AsyncIterator[tuple[str, dict]]:
        usage, pieces, status = None, 0, "cancelled"       # stays so if the client leaves
        try:
            async for chunk in self.provider.stream(prepared.messages, self.settings.llm):
                if chunk.usage:
                    usage = chunk.usage
                else:
                    pieces += 1
                    yield "token", {"text": chunk.text}
            status = "ok"
            yield "done", {"usage": self._usage_out(usage or Usage(0, pieces)),
                           "prompt_version": self.settings.prompt_version}
        except ProviderError as exc:
            status = "error"
            error = to_gateway_error(exc)
            yield "error", {"code": error.code, "message": error.message}
        finally:                                           # also runs when the client disconnects
            if usage is None:                              # no final count: estimate from the words
                usage = Usage(sum(len(m.content.split()) for m in prepared.messages), pieces)
            self._record(prepared, usage, status)

# file: app/main.py
from collections.abc import AsyncIterable
from contextlib import asynccontextmanager
from typing import Annotated

import httpx
from fastapi import APIRouter, Depends, FastAPI, Header, Request
from fastapi.responses import JSONResponse
from fastapi.sse import EventSourceResponse, ServerSentEvent

from app.providers import LLMProvider, Message, PracticeProvider, RetryingProvider
from app.schemas import ChatIn, ChatOut
from app.service import ChatService, GatewayError, Prepared, UsageStore, totals
from app.settings import Settings

router = APIRouter()


def get_settings(request: Request) -> Settings:
    return request.app.state.settings


def get_provider(request: Request) -> LLMProvider:
    return request.app.state.provider


SettingsDep = Annotated[Settings, Depends(get_settings)]
ProviderDep = Annotated[LLMProvider, Depends(get_provider)]


def get_service(request: Request, provider: ProviderDep, settings: SettingsDep) -> ChatService:
    return ChatService(provider, request.app.state.store, settings)


ServiceDep = Annotated[ChatService, Depends(get_service)]


def current_user(settings: SettingsDep, authorization: Annotated[str | None, Header()] = None) -> str:
    name = settings.api_keys.get((authorization or "").removeprefix("Bearer "))
    if name is None:
        raise GatewayError(401, "unauthorized", "Missing or invalid API key", {"WWW-Authenticate": "Bearer"})
    return name


UserDep = Annotated[str, Depends(current_user)]


def admitted(body: ChatIn, user: UserDep, service: ServiceDep) -> Prepared:
    """Runs before the route, so a refusal is a real 429, not an error inside a stream that already said 200."""
    return service.prepare(user, [Message(m.role, m.content) for m in body.messages])


AdmittedDep = Annotated[Prepared, Depends(admitted)]


@router.get("/health")
def health():
    return {"status": "ok"}


@router.get("/health/ready")
async def ready(provider: ProviderDep):
    if not await provider.ping():
        raise GatewayError(503, "provider_unavailable", "The model provider is not reachable")
    return {"status": "ready"}


@router.post("/v1/chat")
async def chat(prepared: AdmittedDep, service: ServiceDep) -> ChatOut:
    return await service.complete(prepared)


@router.post("/v1/chat/stream", response_class=EventSourceResponse)
async def chat_stream(prepared: AdmittedDep, service: ServiceDep) -> AsyncIterable[ServerSentEvent]:
    async for name, data in service.stream(prepared):
        yield ServerSentEvent(data=data, event=name)


@router.get("/v1/usage")
def my_usage(user: UserDep, service: ServiceDep):
    left = service.settings.daily_token_limit - service.store.tokens_today(user)
    return {"user": user, **totals(service.store.today(user)), "tokens_left": max(0, left)}


@router.get("/v1/usage/summary")
def usage_summary(user: UserDep, service: ServiceDep):
    if user != "admin":
        raise GatewayError(403, "forbidden", "Only the admin key can read the summary")
    return service.store.by_prompt_version()


def create_app(settings: Settings | None = None) -> FastAPI:
    @asynccontextmanager
    async def lifespan(app: FastAPI):
        cfg = settings or Settings()                      # fails here, at startup, without a key
        headers = {"Authorization": f"Bearer {cfg.provider_api_key.get_secret_value()}"}
        timeout = httpx.Timeout(cfg.llm.timeout_s, connect=3.0)
        async with httpx.AsyncClient(base_url=cfg.provider_url, headers=headers, timeout=timeout) as client:
            app.state.settings = cfg
            app.state.provider = RetryingProvider(PracticeProvider(client))
            app.state.store = UsageStore()
            yield

    app = FastAPI(title="LLM Gateway", lifespan=lifespan)
    app.include_router(router)

    @app.exception_handler(GatewayError)
    async def handle_gateway_error(request: Request, exc: GatewayError) -> JSONResponse:
        body = {"error": {"code": exc.code, "message": exc.message}}
        return JSONResponse(body, status_code=exc.status, headers=exc.headers)

    return app


app = create_app()

export LLM_PROVIDER_URL=http://localhost:8800
export LLM_PROVIDER_API_KEY=token-nadia-001
export LLM_DAILY_TOKEN_LIMIT=200
uv run fastapi run --port 8000 > server.log 2>&1 &
echo $! > server.pid
curl -s --retry 20 --retry-connrefused --retry-delay 0 http://localhost:8000/health; echo
curl -s http://localhost:8000/health/ready; echo

{"status":"ok"}
{"status":"ready"}

curl -s -X POST http://localhost:8000/v1/chat \
  -H "Authorization: Bearer gw-nadia" -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "user", "content": "What is RAG?"}]}'; echo
curl -s -X POST http://localhost:8000/v1/chat \
  -H "Authorization: Bearer gw-nadia" -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "user", "content": "Write to nadia@example.com about RAG?"}]}'; echo

{"answer":"This is a practice answer about: What is RAG. A real model would say more.","model":"practice-llm-1","prompt_version":"tutor-v2","usage":{"prompt_tokens":18,"completion_tokens":15,"total_tokens":33,"cost_usd":0.000279}}
{"answer":"This is a practice answer about: Write to [EMAIL] about RAG. A real model would say more.","model":"practice-llm-1","prompt_version":"tutor-v2","usage":{"prompt_tokens":20,"completion_tokens":17,"total_tokens":37,"cost_usd":0.000315}}

curl -sN -X POST http://localhost:8000/v1/chat/stream \
  -H "Authorization: Bearer gw-nadia" -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "user", "content": "What is RAG?"}]}' | grep -v '^$' | paste -d' ' - - | sed -n '1,2p;$p'

event: token data: {"text": "This "}
event: token data: {"text": "is "}
event: done data: {"usage": {"prompt_tokens": 18, "completion_tokens": 15, "total_tokens": 33, "cost_usd": 0.000279}, "prompt_version": "tutor-v2"}

grep llm_call server.log | python -c "
import json, sys
for line in sys.stdin:
    r = json.loads(line)
    print(r['user'], r['prompt_version'], r['prompt_tokens'], r['completion_tokens'], r['cost_usd'], r['status'])
"

nadia tutor-v2 18 15 0.000279 ok
nadia tutor-v2 20 17 0.000315 ok
nadia tutor-v2 18 15 0.000279 ok

curl -s http://localhost:8000/v1/usage -H "Authorization: Bearer gw-nadia"; echo
curl -s http://localhost:8000/v1/usage/summary -H "Authorization: Bearer gw-admin"; echo
curl -s http://localhost:8000/v1/usage/summary -H "Authorization: Bearer gw-nadia"; echo

{"user":"nadia","requests":3,"cancelled":0,"prompt_tokens":56,"completion_tokens":47,"cost_usd":0.000873,"tokens_left":97}
{"tutor-v2 / practice-llm-1":{"requests":3,"cancelled":0,"prompt_tokens":56,"completion_tokens":47,"cost_usd":0.000873}}
{"error":{"code":"forbidden","message":"Only the admin key can read the summary"}}

import time
import httpx

GATEWAY = "http://localhost:8000"
NADIA = {"Authorization": "Bearer gw-nadia"}
RAFI = {"Authorization": "Bearer gw-rafi"}
ASK = {"messages": [{"role": "user", "content": "What is RAG?"}]}

with httpx.stream("POST", f"{GATEWAY}/v1/chat/stream", json=ASK, headers=NADIA) as response:
    tokens = 0
    for line in response.iter_lines():
        tokens += line.startswith("event: token")
        if tokens == 3:
            break
print("client left after", tokens, "tokens")

time.sleep(1)                             # the server notices a moment after we leave
usage = httpx.get(f"{GATEWAY}/v1/usage", headers=NADIA).json()
print("requests:", usage["requests"], "| cancelled:", usage["cancelled"])

codes = []
while not codes or codes[-1] == 200:
    codes.append(httpx.post(f"{GATEWAY}/v1/chat", json=ASK, headers=RAFI).status_code)
print(codes)
print(httpx.post(f"{GATEWAY}/v1/chat", json=ASK, headers=RAFI).json())

client left after 3 tokens
requests: 4 | cancelled: 1
[200, 200, 200, 200, 200, 200, 200, 429]
{'error': {'code': 'usage_limit_exceeded', 'message': 'Daily limit of 200 tokens reached'}}

kill $(cat server.pid)
while curl -s -o /dev/null http://localhost:8000/health; do sleep 0.1; done
echo stopped

stopped

# file: tests/test_chat.py
import asyncio

from fastapi.testclient import TestClient

from app.main import create_app, get_provider
from app.providers import Completion, ProviderError, RetryingProvider, StreamChunk, Usage
from app.settings import Settings

QUESTION = {"messages": [{"role": "user", "content": "What is RAG?"}]}
NADIA = {"Authorization": "Bearer gw-nadia"}


class FakeProvider:
    def __init__(self, fail=0, delay=0.0, up=True):
        self.fail, self.delay, self.up, self.calls = fail, delay, up, []

    async def complete(self, messages, config):
        self.calls.append(messages)
        if self.fail > 0:
            self.fail -= 1
            raise ProviderError("overloaded", retryable=True)
        await asyncio.sleep(self.delay)
        return Completion("fake answer", Usage(10, 4))

    async def stream(self, messages, config):
        yield StreamChunk(text="fake ")
        yield StreamChunk(text="answer ")
        yield StreamChunk(usage=Usage(10, 2))

    async def ping(self):
        return self.up


def make_client(provider, **overrides):
    app = create_app(Settings(provider_api_key="test-key", **overrides))
    app.dependency_overrides[get_provider] = lambda: provider
    return TestClient(app)


def test_answer_cost_and_server_owned_prompt():
    fake = FakeProvider()
    with make_client(fake) as client:
        body = client.post("/v1/chat", json=QUESTION, headers=NADIA).json()
    assert body["answer"] == "fake answer"
    assert body["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14, "cost_usd": 0.00009}
    assert fake.calls[0][0].role == "system"


def test_retries_then_succeeds_or_gives_up():
    sleeps = []

    async def fake_sleep(seconds):
        sleeps.append(seconds)

    for failures, expected in ((2, 200), (9, 503)):
        provider = RetryingProvider(FakeProvider(fail=failures), sleep=fake_sleep, jitter=lambda: 1.0)
        with make_client(provider) as client:
            assert client.post("/v1/chat", json=QUESTION, headers=NADIA).status_code == expected
    assert sleeps == [0.5, 1.0, 0.5, 1.0]                 # backoff doubles; three attempts means two pauses


def test_slow_provider_hits_the_deadline():
    with make_client(FakeProvider(delay=1.0), llm={"timeout_s": 0.05}) as client:
        assert client.post("/v1/chat", json=QUESTION, headers=NADIA).status_code == 504


def test_stream_and_readiness():
    with make_client(FakeProvider()) as client:
        text = client.post("/v1/chat/stream", json=QUESTION, headers=NADIA).text
    assert text.count("event: token") == 2 and "event: done" in text
    with make_client(FakeProvider(up=False)) as client:
        assert client.get("/health").status_code == 200
        assert client.get("/health/ready").status_code == 503

uv run python -m pytest -q

....                                                                                         [100%]
4 passed in 1.94s

# file: evals.py
import logging

from fastapi.testclient import TestClient

from app.main import create_app

logging.getLogger("llm").setLevel(logging.WARNING)      # keep the call log out of the report

CASES = [
    ("mentions-topic", "What is RAG?", lambda answer: "RAG" in answer),
    ("hides-email", "I am nadia@example.com. What is RAG?", lambda answer: "nadia@example.com" not in answer),
    ("is-short", "Explain embeddings", lambda answer: len(answer.split()) <= 10),
]

with TestClient(create_app()) as client:
    passed = 0
    for name, question, check in CASES:
        body = {"messages": [{"role": "user", "content": question}]}
        answer = client.post("/v1/chat", json=body, headers={"Authorization": "Bearer gw-nadia"}).json()["answer"]
        ok = check(answer)
        passed += ok
        print("PASS" if ok else "FAIL", name)
    print(f"{passed}/{len(CASES)} passed")

uv run python evals.py

PASS mentions-topic
PASS hides-email
FAIL is-short
2/3 passed

# 1. Easy
from app.redact import redact
print(redact("Call 01712345678. Key: sk-test12345678."))


# 2. Medium (the lookbehind stops it matching inside a longer word or number)
print(redact("Call +8801712345678 or 01712345678."))


# 3. Hard
from tests.test_chat import NADIA, QUESTION, FakeProvider, make_client

with make_client(FakeProvider(), daily_token_limit=1) as client:
    first = client.post("/v1/chat/stream", json=QUESTION, headers=NADIA)
    second = client.post("/v1/chat/stream", json=QUESTION, headers=NADIA)
    usage = client.get("/v1/usage", headers=NADIA).json()
print(first.status_code, second.status_code, second.json()["error"]["code"], usage["requests"])

Call [PHONE]. Key: [KEY].
Call [PHONE] or [PHONE].
200 429 usage_limit_exceeded 1