devplacepy/tests/unit/services/openai_gateway/gateway.py

# retoor <retoor@molodetz.nl>

import json
import requests
from starlette.requests import Request
from tests.conftest import BASE_URL, run_async
from devplacepy.database import get_table, set_setting
from devplacepy.utils import generate_uid
import devplacepy.services.openai_gateway.gateway as gwmod
from devplacepy.services.openai_gateway import GatewayService
class FakeResp_openai_gateway:
    def __init__(self, status=200, payload=None, ctype="application/json", content=b""):
        self.status_code = status
        self._payload = payload
        self.text = json.dumps(payload) if payload is not None else ""
        self.headers = {"content-type": ctype}
        self.content = content

    def json(self):
        if self._payload is None:
            raise ValueError("no json")
        return self._payload
class FakeRequest:
    def __init__(self, method, url, json_body):
        self.method = method
        self.url = url
        self.json_body = json_body
        self.extensions = {}
class FakeClient_openai_gateway:
    def __init__(self, *a, **k):
        self.calls = []

    def build_request(self, method, url, headers=None, json=None, content=None):
        return FakeRequest(method, url, json)

    async def send(self, request):
        self.calls.append((request.url, request.json_body))
        body = request.json_body or {}
        return FakeResp_openai_gateway(
            payload={
                "id": "x",
                "model": body.get("model"),
                "choices": [{"message": {"content": "hi there"}}],
            }
        )

    async def post(self, url, headers=None, json=None, timeout=None):
        self.calls.append((url, json))
        return FakeResp_openai_gateway(
            payload={
                "id": "x",
                "model": json.get("model"),
                "choices": [{"message": {"content": "hi there"}}],
            }
        )

    async def request(self, method, url, headers=None, content=None):
        return FakeResp_openai_gateway(payload={"ok": True})

    async def aclose(self):
        pass
def _make_request_openai_gateway(headers=None, cookies=None):
    headers = headers or {}
    raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()]
    if cookies:
        raw.append(
            (b"cookie", "; ".join(f"{k}={v}" for k, v in cookies.items()).encode())
        )
    return Request(
        {
            "type": "http",
            "method": "POST",
            "path": "/openai/v1/chat/completions",
            "query_string": b"",
            "headers": raw,
            "state": {},
        }
    )
def _make_admin_openai_gateway(role="Admin"):
    username = f"gw_{generate_uid()[:8]}"
    api_key = generate_uid()
    get_table("users").insert(
        {
            "uid": generate_uid(),
            "username": username,
            "terms_version": "1",
            "email": f"{username}@t.dev",
            "api_key": api_key,
            "role": role,
            "is_active": True,
        }
    )
    return username, api_key
def _config(page, **fields):
    page.request.post(f"{BASE_URL}/admin/services/openai/config", form=fields)


def test_model_is_forced(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = True
    cfg["gateway_model"] = "deepseek-chat"
    rt = svc.runtime()
    run_async(
        rt.handle_chat(
            {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "deepseek-chat"


def test_model_route_overrides_upstream(local_db, monkeypatch):
    from devplacepy.services.openai_gateway import routing

    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    routing.provider_store.set(
        routing.ProviderIn(
            name="gwroute",
            base_url="https://routed.example/v1/chat/completions",
            api_key="routed-key",
        )
    )
    routing.model_store.set(
        routing.ModelRouteIn(
            source_model="custom-pro",
            provider="gwroute",
            target_model="vendor/pro",
            kind="chat",
            price_output_per_m=9.0,
        )
    )
    try:
        svc = GatewayService()
        cfg = svc.effective_config()
        cfg["gateway_vision_enabled"] = False
        rt = svc.runtime()
        run_async(
            rt.handle_chat(
                {"model": "custom-pro", "messages": [{"role": "user", "content": "hi"}]},
                cfg,
                ("guest", "test"),
                "test",
                "default",
            )
        )
        url, body = rt._client.calls[-1]
        assert str(url) == "https://routed.example/v1/chat/completions"
        assert body["model"] == "vendor/pro"
    finally:
        routing.model_store.remove("custom-pro")
        routing.provider_store.remove("gwroute")


def test_vision_rewrites_image_to_text(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_vision_enabled"] = True
    cfg["gateway_vision_key"] = ""  # no key -> placeholder text, still rewrites to str
    rt = svc.runtime()
    msgs = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "what is this"},
                {"type": "image_url", "image_url": {"url": "http://x/y.png"}},
            ],
        }
    ]
    run_async(rt.handle_chat({"messages": msgs}, cfg, ("guest", "test"), "test", "default"))
    sent = rt._client.calls[-1][1]["messages"][0]["content"]
    assert isinstance(sent, str) and "vision" in sent.lower()


def test_streaming_emits_sse(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    rt = svc.runtime()
    resp = run_async(
        rt.handle_chat(
            {"messages": [{"role": "user", "content": "hi"}], "stream": True},
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )

    async def drain():
        out = []
        async for chunk in resp.body_iterator:
            out.append(chunk if isinstance(chunk, str) else chunk.decode())
        return "".join(out)

    body = run_async(drain())
    assert "chat.completion.chunk" in body
    assert "[DONE]" in body


def test_authorize_static_access_key(local_db):
    set_setting("gateway_require_auth", "1")
    set_setting("gateway_access_key", "topsecret")
    try:
        svc = GatewayService()
        assert svc.authorize(_make_request_openai_gateway(headers={"X-API-KEY": "topsecret"})) is True
        assert svc.authorize(_make_request_openai_gateway(headers={"X-API-KEY": "wrong"})) is False
    finally:
        set_setting("gateway_access_key", "")
        set_setting("gateway_require_auth", "1")


def test_authorize_admin_and_user_toggles(local_db):
    set_setting("gateway_require_auth", "1")
    set_setting("gateway_access_key", "")
    set_setting("gateway_allow_admins", "1")
    set_setting("gateway_allow_users", "0")
    try:
        _, admin_key = _make_admin_openai_gateway(role="Admin")
        _, member_key = _make_admin_openai_gateway(role="Member")
        svc = GatewayService()

        assert (
            svc.authorize(_make_request_openai_gateway(headers={"Authorization": f"Bearer {admin_key}"}))
            is True
        )
        assert svc.authorize(_make_request_openai_gateway(headers={"X-API-KEY": member_key})) is False

        set_setting("gateway_allow_users", "1")
        assert svc.authorize(_make_request_openai_gateway(headers={"X-API-KEY": member_key})) is True
    finally:
        set_setting("gateway_allow_admins", "1")
        set_setting("gateway_allow_users", "1")
        set_setting("gateway_require_auth", "1")
        set_setting("gateway_access_key", "")


def test_authorize_require_auth_off_is_open(local_db):
    set_setting("gateway_require_auth", "0")
    try:
        svc = GatewayService()
        assert svc.authorize(_make_request_openai_gateway()) is True
    finally:
        set_setting("gateway_require_auth", "1")


class FakeEmbedClient_openai_gateway:
    def __init__(self, *a, **k):
        self.calls = []

    def build_request(self, method, url, headers=None, json=None, content=None):
        return FakeRequest(method, url, json)

    async def send(self, request):
        self.calls.append((request.url, request.json_body))
        body = request.json_body or {}
        return FakeResp_openai_gateway(
            payload={
                "object": "list",
                "model": body.get("model"),
                "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
                "usage": {"prompt_tokens": 5, "total_tokens": 5},
            }
        )

    async def aclose(self):
        pass


def test_embeddings_remaps_alias_to_configured_model(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeEmbedClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = False
    cfg["gateway_embed_enabled"] = True
    cfg["gateway_embed_model"] = "qwen/qwen3-embedding-8b"
    rt = svc.runtime()
    run_async(
        rt.handle_embeddings(
            {"model": "molodetz~embed", "input": "hello"},
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "qwen/qwen3-embedding-8b"


def test_embeddings_success_records_ledger(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeEmbedClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = True
    cfg["gateway_embed_enabled"] = True
    cfg["gateway_embed_model"] = "qwen/qwen3-embedding-8b"
    rt = svc.runtime()
    before = rt.embed_calls
    resp = run_async(
        rt.handle_embeddings(
            {"model": "molodetz~embed", "input": "hello"},
            cfg,
            ("guest", "ledger_probe"),
            "test",
            "default",
        )
    )
    assert resp.status_code == 200
    payload = json.loads(bytes(resp.body).decode())
    assert payload["data"][0]["embedding"] == [0.1, 0.2]
    assert rt.embed_calls == before + 1
    row = get_table("gateway_usage_ledger").find_one(owner_id="ledger_probe")
    assert row is not None
    assert row["backend"] == "embed"
    assert row["endpoint"] == "embeddings"
    assert row["requested_model"] == "molodetz~embed"
    assert row["model"] == "qwen/qwen3-embedding-8b"
    assert row["success"] == 1


def test_embeddings_disabled_returns_503(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeEmbedClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_embed_enabled"] = False
    rt = svc.runtime()
    resp = run_async(
        rt.handle_embeddings(
            {"input": "hello"}, cfg, ("guest", "test"), "test", "default"
        )
    )
    assert resp.status_code == 503


class FakeImageClient_openai_gateway:
    def __init__(self, *a, **k):
        self.calls = []

    def build_request(self, method, url, headers=None, json=None, content=None):
        return FakeRequest(method, url, json)

    async def send(self, request):
        self.calls.append((request.url, request.json_body))
        body = request.json_body or {}
        return FakeResp_openai_gateway(
            payload={
                "created": 1,
                "model": body.get("model"),
                "data": [{"b64_json": "aGVsbG8="}],
                "usage": {"cost": 0.05},
            }
        )

    async def aclose(self):
        pass


def test_images_remaps_alias_to_configured_model(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = False
    cfg["gateway_image_enabled"] = True
    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
    rt = svc.runtime()
    run_async(
        rt.handle_images(
            {
                "model": "molodetz-img-small",
                "prompt": "award emblem",
                "size": "512x512",
            },
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "black-forest-labs/flux.2-pro"


def test_image_route_overrides_upstream(local_db, monkeypatch):
    from devplacepy.services.openai_gateway import routing

    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    routing.provider_store.set(
        routing.ProviderIn(
            name="gwimg",
            base_url="https://routed.example/v1/chat/completions",
            api_key="img-key",
        )
    )
    routing.model_store.set(
        routing.ModelRouteIn(
            source_model="custom-img",
            provider="gwimg",
            target_model="vendor/flux-pro",
            kind="image",
            price_input_per_m=0.06,
        )
    )
    try:
        svc = GatewayService()
        cfg = svc.effective_config()
        rt = svc.runtime()
        run_async(
            rt.handle_images(
                {
                    "model": "custom-img",
                    "prompt": "trophy",
                    "response_format": "b64_json",
                },
                cfg,
                ("guest", "test"),
                "test",
                "default",
            )
        )
        url, body = rt._client.calls[-1]
        assert str(url) == "https://routed.example/v1/images"
        assert body["model"] == "vendor/flux-pro"
    finally:
        routing.model_store.remove("custom-img")
        routing.provider_store.remove("gwimg")


def test_images_success_records_ledger(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_image_enabled"] = True
    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
    rt = svc.runtime()
    before = rt.image_calls
    resp = run_async(
        rt.handle_images(
            {"model": "molodetz-img-small", "prompt": "badge"},
            cfg,
            ("guest", "img_ledger"),
            "test",
            "default",
        )
    )
    assert resp.status_code == 200
    payload = json.loads(bytes(resp.body).decode())
    assert payload["data"][0]["b64_json"] == "aGVsbG8="
    assert rt.image_calls == before + 1
    row = get_table("gateway_usage_ledger").find_one(owner_id="img_ledger")
    assert row is not None
    assert row["backend"] == "image"
    assert row["endpoint"] == "images/generations"
    assert row["requested_model"] == "molodetz-img-small"
    assert row["model"] == "black-forest-labs/flux.2-pro"
    assert row["success"] == 1
    assert float(row["cost_usd"]) == 0.05


def test_images_disabled_returns_503(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_image_enabled"] = False
    rt = svc.runtime()
    resp = run_async(
        rt.handle_images(
            {"prompt": "badge"}, cfg, ("guest", "test"), "test", "default"
        )
    )
    assert resp.status_code == 503


def test_compute_cost_embed_branch():
    from devplacepy.services.openai_gateway.usage import (
        Pricing,
        compute_cost,
        normalize_usage,
    )

    pricing = Pricing(
        chat_cache_hit_per_m=0.0,
        chat_cache_miss_per_m=0.0,
        chat_output_per_m=0.0,
        vision_input_per_m=0.0,
        vision_output_per_m=0.0,
        embed_input_per_m=0.01,
    )
    usage = {"prompt_tokens": 1_000_000, "total_tokens": 1_000_000}
    norm = normalize_usage(usage)
    total, input_cost, output_cost, native = compute_cost(
        usage, norm, pricing, "embed"
    )
    assert native is False
    assert output_cost == 0.0
    assert abs(total - 0.01) < 1e-9
    assert abs(input_cost - 0.01) < 1e-9


def test_app_reference_stored_in_ledger_for_chat(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = True
    cfg["gateway_model"] = "deepseek-chat"
    rt = svc.runtime()
    run_async(
        rt.handle_chat(
            {"messages": [{"role": "user", "content": "ref_test"}]},
            cfg,
            ("guest", "ref_probe"),
            "test-ua",
            "my-custom-app",
        )
    )
    row = get_table("gateway_usage_ledger").find_one(owner_id="ref_probe")
    assert row is not None
    assert row["app_reference"] == "my-custom-app"
    assert row["backend"] == "chat"


def test_app_reference_defaults_when_not_passed(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = True
    cfg["gateway_model"] = "deepseek-chat"
    rt = svc.runtime()
    run_async(
        rt.handle_chat(
            {"messages": [{"role": "user", "content": "default_test"}]},
            cfg,
            ("guest", "default_probe"),
            "test-ua",
            "default",
        )
    )
    row = get_table("gateway_usage_ledger").find_one(owner_id="default_probe")
    assert row is not None
    assert row["app_reference"] == "default"


def test_app_reference_stored_in_ledger_for_embeddings(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeEmbedClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_embed_enabled"] = True
    cfg["gateway_embed_model"] = "qwen/qwen3-embedding-8b"
    rt = svc.runtime()
    run_async(
        rt.handle_embeddings(
            {"input": "hello"},
            cfg,
            ("guest", "embed_ref"),
            "test-ua",
            "devplace-embed-test-v-1-0-0",
        )
    )
    row = get_table("gateway_usage_ledger").find_one(owner_id="embed_ref")
    assert row is not None
    assert row["app_reference"] == "devplace-embed-test-v-1-0-0"
    assert row["backend"] == "embed"


def test_app_reference_stored_in_ledger_for_images(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_image_enabled"] = True
    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
    rt = svc.runtime()
    run_async(
        rt.handle_images(
            {"prompt": "ref badge"},
            cfg,
            ("guest", "img_ref"),
            "test-ua",
            "devplace-image-test-v-1-0-0",
        )
    )
    row = get_table("gateway_usage_ledger").find_one(owner_id="img_ref")
    assert row is not None
    assert row["app_reference"] == "devplace-image-test-v-1-0-0"
    assert row["backend"] == "image"


def test_app_reference_column_exists(local_db):
    db = get_table("gateway_usage_ledger").db
    if "gateway_usage_ledger" not in db.tables:
        return
    rows = list(db.query("PRAGMA table_info('gateway_usage_ledger')"))
    column_names = [r["name"] for r in rows]
    assert "app_reference" in column_names


def test_chat_unknown_model_falls_back_to_default(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = False
    cfg["gateway_model"] = "deepseek-v4-flash"
    rt = svc.runtime()
    run_async(
        rt.handle_chat(
            {"model": "nonexistent-model-v99", "messages": [{"role": "user", "content": "hi"}]},
            cfg,
            ("guest", "fallback_test_chat"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "deepseek-v4-flash"


def test_embeddings_unknown_model_falls_back_to_default(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeEmbedClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = False
    cfg["gateway_embed_enabled"] = True
    cfg["gateway_embed_model"] = "qwen/qwen3-embedding-8b"
    rt = svc.runtime()
    run_async(
        rt.handle_embeddings(
            {"model": "nonexistent-embed-model", "input": "hello"},
            cfg,
            ("guest", "fallback_test_embed"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "qwen/qwen3-embedding-8b"


def test_images_unknown_model_falls_back_to_default(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeImageClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    cfg["gateway_force_model"] = False
    cfg["gateway_image_enabled"] = True
    cfg["gateway_image_model"] = "black-forest-labs/flux.2-pro"
    rt = svc.runtime()
    run_async(
        rt.handle_images(
            {"model": "nonexistent-image-model", "prompt": "test badge"},
            cfg,
            ("guest", "fallback_test_img"),
            "test",
            "default",
        )
    )
    assert rt._client.calls[-1][1]["model"] == "black-forest-labs/flux.2-pro"


class FakeClientWithUsage_openai_gateway(FakeClient_openai_gateway):
    async def send(self, request):
        self.calls.append((request.url, request.json_body))
        body = request.json_body or {}
        return FakeResp_openai_gateway(
            payload={
                "id": "x",
                "model": body.get("model"),
                "choices": [{"message": {"content": "hi there"}}],
                "usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
            }
        )


def test_stream_options_stripped_from_upstream(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClient_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    rt = svc.runtime()
    run_async(
        rt.handle_chat(
            {
                "messages": [{"role": "user", "content": "hi"}],
                "stream": True,
                "stream_options": {"include_usage": True},
            },
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )
    upstream_payload = rt._client.calls[-1][1]
    assert "stream_options" not in upstream_payload
    assert upstream_payload["stream"] is False


def test_include_usage_emits_usage_chunk(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClientWithUsage_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    rt = svc.runtime()
    resp = run_async(
        rt.handle_chat(
            {
                "messages": [{"role": "user", "content": "hi"}],
                "stream": True,
                "stream_options": {"include_usage": True},
            },
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )

    async def drain():
        out = []
        async for chunk in resp.body_iterator:
            out.append(chunk if isinstance(chunk, str) else chunk.decode())
        return "".join(out)

    body = run_async(drain())
    assert '"usage"' in body
    assert '"total_tokens": 10' in body
    assert "[DONE]" in body


def test_no_usage_chunk_without_include_usage(local_db, monkeypatch):
    monkeypatch.setattr(gwmod.httpx, "AsyncClient", FakeClientWithUsage_openai_gateway)
    svc = GatewayService()
    cfg = svc.effective_config()
    rt = svc.runtime()
    resp = run_async(
        rt.handle_chat(
            {"messages": [{"role": "user", "content": "hi"}], "stream": True},
            cfg,
            ("guest", "test"),
            "test",
            "default",
        )
    )

    async def drain():
        out = []
        async for chunk in resp.body_iterator:
            out.append(chunk if isinstance(chunk, str) else chunk.decode())
        return "".join(out)

    body = run_async(drain())
    assert '"usage"' not in body
    assert "[DONE]" in body


def test_models_endpoint_publishes_molodetz(local_db):
    from devplacepy.services.openai_gateway import routing

    routing.seed_default_deepseek_routes()
    try:
        svc = GatewayService()
        resp = svc._models_response()
        payload = json.loads(resp.body.decode())
        ids = {m["id"] for m in payload["data"]}
        assert payload["object"] == "list"
        assert {"molodetz", "molodetz-pro"} <= ids
        for model in payload["data"]:
            assert model["object"] == "model"
            assert model["owned_by"] == "molodetz"
    finally:
        routing.model_store.remove("molodetz")
        routing.model_store.remove("molodetz-pro")