pablan/backend/tests/fake_openai.py

"""Minimal fake OpenAI-compatible server as an ASGI app (no real LLM).

Wire it into the client via httpx.ASGITransport — see the fake_llm fixture.
Chat behavior is scripted through `chat_responses`; each entry is one of:
  {"content": "..."}          plain completion content
  {"chunks": ["a", "b"]}      streamed delta pieces
  {"status": 500}             induced HTTP error
An empty script falls back to `default_content`.
"""

import json
from dataclasses import dataclass, field
from typing import Any

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse

_USAGE = {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}


@dataclass
class FakeOpenAI:
    chat_responses: list[dict[str, Any]] = field(default_factory=list)
    default_content: str = "pong"
    embedding_dim: int = 8
    # What GET /v1/models reports; None means the route is missing (404).
    served_models: list[str] | None = field(
        default_factory=lambda: ["gemma-3-27b", "bge-m3"]
    )
    requests: list[dict[str, Any]] = field(default_factory=list)

    def __post_init__(self) -> None:
        self.app = self._build_app()

    def _build_app(self) -> FastAPI:
        app = FastAPI()

        @app.post("/v1/chat/completions")
        async def chat_completions(request: Request) -> Any:
            body = await request.json()
            self.requests.append(body)
            plan = (
                self.chat_responses.pop(0)
                if self.chat_responses
                else {"content": self.default_content}
            )
            if "status" in plan:
                return JSONResponse(
                    {"error": {"message": "induced failure", "type": "server_error"}},
                    status_code=plan["status"],
                )
            content = plan.get("content", self.default_content)
            if body.get("stream"):
                pieces = plan.get("chunks", [content])
                include_usage = bool(
                    (body.get("stream_options") or {}).get("include_usage")
                )
                return StreamingResponse(
                    self._stream(body["model"], pieces, include_usage),
                    media_type="text/event-stream",
                )
            return JSONResponse(
                {
                    "id": "cmpl-fake",
                    "object": "chat.completion",
                    "created": 0,
                    "model": body["model"],
                    "choices": [
                        {
                            "index": 0,
                            "message": {"role": "assistant", "content": content},
                            "finish_reason": "stop",
                        }
                    ],
                    "usage": _USAGE,
                }
            )

        @app.post("/v1/embeddings")
        async def embeddings(request: Request) -> Any:
            body = await request.json()
            self.requests.append(body)
            inputs = body["input"]
            if isinstance(inputs, str):
                inputs = [inputs]
            return JSONResponse(
                {
                    "object": "list",
                    "model": body["model"],
                    "data": [
                        {
                            "object": "embedding",
                            "index": i,
                            "embedding": [float(i)] * self.embedding_dim,
                        }
                        for i in range(len(inputs))
                    ],
                    "usage": {
                        "prompt_tokens": len(inputs),
                        "total_tokens": len(inputs),
                    },
                }
            )

        @app.get("/v1/models")
        async def models() -> Any:
            # `served_models = None` mimics an endpoint without the route,
            # which plenty of OpenAI-compatible servers genuinely lack.
            if self.served_models is None:
                return JSONResponse(
                    {"error": {"message": "not found", "type": "invalid_request"}},
                    status_code=404,
                )
            return JSONResponse(
                {
                    "object": "list",
                    "data": [
                        {"id": name, "object": "model", "owned_by": "fake"}
                        for name in self.served_models
                    ],
                }
            )

        return app

    @staticmethod
    async def _stream(model: str, pieces: list[str], include_usage: bool) -> Any:
        def chunk(delta: dict[str, Any], finish: str | None) -> str:
            payload = {
                "id": "cmpl-fake",
                "object": "chat.completion.chunk",
                "created": 0,
                "model": model,
                "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
            }
            return f"data: {json.dumps(payload)}\n\n"

        for piece in pieces:
            yield chunk({"content": piece}, None)
        yield chunk({}, "stop")
        if include_usage:
            final = {
                "id": "cmpl-fake",
                "object": "chat.completion.chunk",
                "created": 0,
                "model": model,
                "choices": [],
                "usage": _USAGE,
            }
            yield f"data: {json.dumps(final)}\n\n"
        yield "data: [DONE]\n\n"