DocsGPT/scripts/e2e/mock_llm.py

"""OpenAI-compatible stub server for the DocsGPT e2e test suite.

Speaks the minimum subset of the OpenAI HTTP API that DocsGPT's ``openai``
Python client needs:

* ``POST /v1/chat/completions`` (streaming + non-streaming, tool calls via fixture)
* ``POST /v1/embeddings`` (deterministic hash-seeded vectors)
* ``GET /healthz`` (liveness probe for ``scripts/e2e/up.sh``)

The server is **deterministic**: the same request always returns the same
response. Requests are fingerprinted by SHA-256 of a canonical JSON encoding
of ``(model, messages, tool_choice)``. If a fixture file matching that hash
exists under ``mock_llm_fixtures/<hash>.json`` it wins; otherwise a generic
"I don't know" fallback is returned and the hash + request is logged to stderr
so a developer can promote it into a fixture later.

Run standalone (does NOT import anything from ``application/``). Python 3.11+.
Flask is the only non-stdlib dependency and is already in
``application/requirements.txt``.

Usage::

    python scripts/e2e/mock_llm.py

Defaults to ``127.0.0.1:7899`` to match the ``OPENAI_BASE_URL`` referenced in
``e2e-plan.md`` Appendix A.
"""

from __future__ import annotations

import hashlib
import json
import os
import random
import sys
import time
from pathlib import Path
from typing import Any

from flask import Flask, Response, jsonify, request, stream_with_context

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

HOST = os.environ.get("MOCK_LLM_HOST", "127.0.0.1")
PORT = int(os.environ.get("MOCK_LLM_PORT", "7899"))
FIXTURES_DIR = Path(__file__).parent / "mock_llm_fixtures"
EMBEDDING_DIM = 768
GENERIC_FALLBACK_TEXT = (
    "I don't have enough information to answer that from the provided sources."
)
STREAM_CHUNK_COUNT = 5

app = Flask(__name__)


# ---------------------------------------------------------------------------
# CORS — permissive; stub trusts its port
# ---------------------------------------------------------------------------


@app.after_request
def _add_cors_headers(response: Response) -> Response:
    response.headers["Access-Control-Allow-Origin"] = "*"
    response.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "*"
    return response


@app.route("/v1/chat/completions", methods=["OPTIONS"])
@app.route("/v1/embeddings", methods=["OPTIONS"])
def _cors_preflight() -> Response:
    return Response(status=204)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _canonical_messages(messages: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
    """Return a minimal, stable representation of the messages array.

    We keep only fields that are semantically meaningful for fingerprinting a
    request. Extra keys from the OpenAI client (e.g. ``name``, ``tool_call_id``)
    are preserved because they *do* change the intended response.
    """

    if not messages:
        return []
    out: list[dict[str, Any]] = []
    for msg in messages:
        if not isinstance(msg, dict):
            continue
        # Content may be a string OR a list of content-part dicts (vision / tool).
        # Serialize both forms deterministically.
        entry: dict[str, Any] = {
            "role": msg.get("role"),
            "content": msg.get("content"),
        }
        for key in ("name", "tool_call_id", "tool_calls"):
            if key in msg:
                entry[key] = msg[key]
        out.append(entry)
    return out


def _compute_request_digest(payload: dict[str, Any]) -> str:
    """SHA-256 fingerprint of ``(model, messages, tool_choice)``.

    Kept narrow on purpose — temperature / top_p / seed / max_tokens should
    NOT influence which canned answer we return; those are knobs the app may
    flap on across runs.
    """

    canonical = {
        "model": payload.get("model"),
        "messages": _canonical_messages(payload.get("messages")),
        "tool_choice": payload.get("tool_choice"),
    }
    blob = json.dumps(canonical, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()


def _load_fixture(digest: str) -> dict[str, Any] | None:
    """Return the parsed fixture dict for ``digest``, or ``None`` if missing/bad."""

    path = FIXTURES_DIR / f"{digest}.json"
    if not path.is_file():
        return None
    try:
        with path.open("r", encoding="utf-8") as fh:
            data = json.load(fh)
    except (OSError, json.JSONDecodeError) as exc:
        sys.stderr.write(f"[mock-llm] failed to load fixture {path}: {exc}\n")
        sys.stderr.flush()
        return None
    return data


def _estimate_tokens(text: str) -> int:
    """Rough 4-chars-per-token estimate (OpenAI's own ballpark)."""

    if not text:
        return 0
    return max(1, len(text) // 4)


def _messages_text(messages: list[dict[str, Any]] | None) -> str:
    """Concatenate message contents for prompt-token estimation."""

    if not messages:
        return ""
    parts: list[str] = []
    for msg in messages:
        if not isinstance(msg, dict):
            continue
        content = msg.get("content")
        if isinstance(content, str):
            parts.append(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict):
                    text = part.get("text")
                    if isinstance(text, str):
                        parts.append(text)
    return "\n".join(parts)


def _split_into_chunks(text: str, count: int) -> list[str]:
    """Split ``text`` into roughly ``count`` pieces by character length.

    Guarantees at least one chunk even for the empty string (so streaming
    clients still see a delta before ``[DONE]``).
    """

    if count <= 0:
        return [text]
    if not text:
        return [""]
    n = len(text)
    size = max(1, (n + count - 1) // count)
    chunks = [text[i : i + size] for i in range(0, n, size)]
    if not chunks:
        chunks = [""]
    return chunks


# ---------------------------------------------------------------------------
# Chat completions
# ---------------------------------------------------------------------------


def _resolve_chat_response(
    payload: dict[str, Any], digest: str
) -> tuple[str, list[dict[str, Any]] | None, str, dict[str, int]]:
    """Return ``(content, tool_calls, finish_reason, usage)`` for ``payload``.

    Looks up a fixture by digest first; falls back to the generic response if
    no fixture is present, and logs the miss so the dev can convert it.
    """

    fixture = _load_fixture(digest)
    if fixture is None:
        sys.stderr.write(f"[mock-llm] unknown fixture hash {digest}\n")
        try:
            sys.stderr.write(
                "[mock-llm] request: "
                + json.dumps(payload, sort_keys=True, ensure_ascii=False)
                + "\n"
            )
        except (TypeError, ValueError):
            sys.stderr.write("[mock-llm] request: <unserializable>\n")
        sys.stderr.flush()
        content = GENERIC_FALLBACK_TEXT
        tool_calls: list[dict[str, Any]] | None = None
        finish_reason = "stop"
        prompt_tokens = _estimate_tokens(_messages_text(payload.get("messages")))
        completion_tokens = _estimate_tokens(content)
        usage = {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        }
        return content, tool_calls, finish_reason, usage

    response = fixture.get("response") or {}
    content = response.get("content") or ""
    tool_calls = response.get("tool_calls")
    finish_reason = response.get("finish_reason") or "stop"
    fixture_usage = response.get("usage") or {}
    prompt_tokens = int(
        fixture_usage.get(
            "prompt_tokens",
            _estimate_tokens(_messages_text(payload.get("messages"))),
        )
    )
    completion_tokens = int(
        fixture_usage.get("completion_tokens", _estimate_tokens(content))
    )
    usage = {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
    }
    return content, tool_calls, finish_reason, usage


def _chat_completion_envelope(
    *,
    digest: str,
    model: str,
    content: str,
    tool_calls: list[dict[str, Any]] | None,
    finish_reason: str,
    usage: dict[str, int],
) -> dict[str, Any]:
    message: dict[str, Any] = {"role": "assistant", "content": content}
    if tool_calls:
        message["tool_calls"] = tool_calls
    return {
        "id": f"chatcmpl-e2e-{digest[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": message,
                "finish_reason": finish_reason,
            }
        ],
        "usage": usage,
    }


def _sse(payload: dict[str, Any]) -> str:
    return "data: " + json.dumps(payload, ensure_ascii=False) + "\n\n"


def _stream_chat_response(
    *,
    digest: str,
    model: str,
    content: str,
    tool_calls: list[dict[str, Any]] | None,
    finish_reason: str,
):
    """Generator yielding SSE frames that match the OpenAI streaming protocol."""

    created = int(time.time())
    completion_id = f"chatcmpl-e2e-{digest[:12]}"

    def _base_chunk(delta: dict[str, Any], final: bool = False) -> dict[str, Any]:
        return {
            "id": completion_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": delta,
                    "finish_reason": finish_reason if final else None,
                }
            ],
        }

    # Opening role delta — matches OpenAI's real behavior.
    yield _sse(_base_chunk({"role": "assistant", "content": ""}))

    if tool_calls:
        # Emit tool calls in one delta; content streaming is skipped when
        # tool_calls are present, matching what RAG code paths expect.
        yield _sse(_base_chunk({"tool_calls": tool_calls}))
        yield _sse(_base_chunk({}, final=True))
    else:
        chunks = _split_into_chunks(content, STREAM_CHUNK_COUNT)
        last_index = len(chunks) - 1
        for i, piece in enumerate(chunks):
            yield _sse(_base_chunk({"content": piece}, final=(i == last_index)))

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


@app.post("/v1/chat/completions")
def chat_completions() -> Response:
    payload = request.get_json(silent=True) or {}
    model = payload.get("model") or "gpt-4o-mini"
    stream = bool(payload.get("stream"))

    digest = _compute_request_digest(payload)
    content, tool_calls, finish_reason, usage = _resolve_chat_response(payload, digest)

    if stream:
        generator = _stream_chat_response(
            digest=digest,
            model=model,
            content=content,
            tool_calls=tool_calls,
            finish_reason=finish_reason,
        )
        response = Response(
            stream_with_context(generator),
            mimetype="text/event-stream",
        )
        response.headers["Cache-Control"] = "no-cache"
        response.headers["X-Accel-Buffering"] = "no"
        # Flask strips Content-Length on streamed responses; make sure we don't
        # accidentally set one. Nothing to do here — just documenting.
        return response

    envelope = _chat_completion_envelope(
        digest=digest,
        model=model,
        content=content,
        tool_calls=tool_calls,
        finish_reason=finish_reason,
        usage=usage,
    )
    return jsonify(envelope)


# ---------------------------------------------------------------------------
# Embeddings
# ---------------------------------------------------------------------------


def _deterministic_embedding(text: str) -> list[float]:
    """Hash-seeded 768-dim float vector in [-1, 1).

    Never all-zero: seeded RNG on a non-trivial hash of ``text`` plus a small
    non-zero offset so degenerate vector-store checks pass even if
    ``text`` itself is empty.
    """

    seed = int(hashlib.sha256(text.encode("utf-8")).hexdigest()[:16], 16) & 0xFFFFFFFF
    rng = random.Random(seed)
    vec = [rng.uniform(-1.0, 1.0) for _ in range(EMBEDDING_DIM)]
    # Guarantee non-degeneracy: nudge the first component away from 0 if the
    # seeded draw happens to produce a very small value.
    if abs(vec[0]) < 1e-6:
        vec[0] = 0.1
    return vec


@app.post("/v1/embeddings")
@app.post("/v1/v1/embeddings")
def embeddings() -> Response:
    payload = request.get_json(silent=True) or {}
    model = payload.get("model") or "text-embedding-3-small"
    raw_input = payload.get("input", "")

    if isinstance(raw_input, str):
        inputs: list[str] = [raw_input]
    elif isinstance(raw_input, list):
        inputs = [str(item) if not isinstance(item, str) else item for item in raw_input]
    else:
        inputs = [str(raw_input)]

    data = [
        {
            "object": "embedding",
            "index": i,
            "embedding": _deterministic_embedding(text),
        }
        for i, text in enumerate(inputs)
    ]
    total_tokens = sum(_estimate_tokens(text) for text in inputs)
    return jsonify(
        {
            "object": "list",
            "data": data,
            "model": model,
            "usage": {
                "prompt_tokens": total_tokens,
                "total_tokens": total_tokens,
            },
        }
    )


# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------


@app.get("/healthz")
def healthz() -> Response:
    return jsonify({"ok": True})


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------


def main() -> None:
    FIXTURES_DIR.mkdir(parents=True, exist_ok=True)
    sys.stderr.write(
        f"[mock-llm] listening on http://{HOST}:{PORT} "
        f"(fixtures: {FIXTURES_DIR})\n"
    )
    sys.stderr.flush()
    # threaded=True so that concurrent streaming + embeddings requests from
    # the Flask backend + Celery worker don't serialize behind each other.
    app.run(host=HOST, port=PORT, debug=False, use_reloader=False, threaded=True)


if __name__ == "__main__":
    main()