File size: 13,716 Bytes

730c5bb

"""bankml operator backend: registry, env, receipts, typed refusals, SSE streaming, wiring.

No network: every HTTP exchange goes through httpx.MockTransport.
"""

from __future__ import annotations

import asyncio
import json

import httpx
import pytest
from fastapi.testclient import TestClient

from mindxtrain.governance import panel as P
from mindxtrain.models.registry import ChatMessage, ChatRequest, build_backend, list_backends
from mindxtrain.operator import app as operator_app
from mindxtrain.operator.backends.bankml import (
    DEFAULT_BANKML_BASE_URL,
    BankmlBackend,
    BankmlError,
    BankmlRefusal,
    bankml_root_url,
)

_REAL_BANKML_REACHABLE = operator_app._bankml_reachable  # captured before conftest patches it

RECEIPT = {
    "bankml": "0.3.4",
    "engine": "native",
    "model_sha256": "6b64c748d96ad26fd72402299bd27b2ae82f489bd0469498dd18eb6054058266",
    "guard": "play",
    "prompt_tokens": 12,
    "completion_tokens": 3,
    "ttft_ms": 40,
    "wall_ms": 90,
    "response_sha256": "a" * 64,
    "request_sha256": "b" * 64,
    "signed": False,
}


def _req(stream: bool = False) -> ChatRequest:
    return ChatRequest(
        model="mindx-gen39",
        messages=[ChatMessage(role="user", content="Who are you?")],
        temperature=0.0,
        max_tokens=16,
        stream=stream,
    )


def _completion() -> dict[str, object]:
    return {
        "choices": [{"index": 0, "message": {"role": "assistant", "content": "I am mindX."},
                     "finish_reason": "stop"}],
        "model": "mindx-gen39-F16.gguf",
        "object": "chat.completion",
        "usage": {"completion_tokens": 3, "prompt_tokens": 12, "total_tokens": 15},
        "bankml_receipt": RECEIPT,
    }


# ---- registry + env -----------------------------------------------------------


def test_registered_and_built_by_name() -> None:
    assert "bankml" in list_backends()
    assert isinstance(build_backend("bankml"), BankmlBackend)


def test_default_and_env_base_url(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("MINDXTRAIN_BANKML_BASE_URL", raising=False)
    assert BankmlBackend().base_url == DEFAULT_BANKML_BASE_URL
    monkeypatch.setenv("MINDXTRAIN_BANKML_BASE_URL", "http://127.0.0.1:9999/v1/")
    b = BankmlBackend()
    assert b.base_url == "http://127.0.0.1:9999/v1"
    assert bankml_root_url() == "http://127.0.0.1:9999"


def test_never_sends_an_openai_key(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("MINDXTRAIN_OPENAI_API_KEY", "sk-should-not-leak")
    assert "Authorization" not in BankmlBackend()._headers()


# ---- chat ---------------------------------------------------------------------


def test_chat_keeps_the_receipt() -> None:
    seen: dict[str, object] = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen["url"] = str(request.url)
        seen["body"] = json.loads(request.content)
        return httpx.Response(200, json=_completion())

    b = BankmlBackend(seed=7, transport=httpx.MockTransport(handler))
    resp = asyncio.run(b.chat(_req()))
    assert seen["url"] == "http://127.0.0.1:18093/v1/chat/completions"
    body = seen["body"]
    assert isinstance(body, dict)
    assert body["seed"] == 7 and body["temperature"] == 0.0 and body["stream"] is False
    assert resp.content == "I am mindX."
    assert resp.prompt_tokens == 12 and resp.completion_tokens == 3
    assert resp.receipt == RECEIPT
    assert b.last_receipt == RECEIPT


def test_chat_400_is_a_typed_refusal_and_is_not_retried() -> None:
    calls = {"n": 0}

    def handler(request: httpx.Request) -> httpx.Response:
        calls["n"] += 1
        return httpx.Response(400, text="mirostat: not reproduced; bankML reproduces llama.cpp's "
                                        "top-k, top-p, min-p and temperature")

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    with pytest.raises(BankmlRefusal) as info:
        asyncio.run(b.chat(_req()))
    assert info.value.status_code == 400
    assert info.value.reason.startswith("mirostat: not reproduced")
    assert "bankml refused (400)" in str(info.value)
    assert calls["n"] == 1


def test_ollama_style_json_error_body_is_read() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(400, json={"error": "tools: tool calling ... not reproduced yet"})

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    with pytest.raises(BankmlRefusal) as info:
        asyncio.run(b.chat(_req()))
    assert info.value.reason.startswith("tools:")


def test_non_400_errors_are_bankml_errors_not_refusals() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(503, text="model changed since bankml verified it")

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    with pytest.raises(BankmlError) as info:
        asyncio.run(b.chat(_req()))
    assert not isinstance(info.value, BankmlRefusal)
    assert info.value.status_code == 503


# ---- stream -------------------------------------------------------------------


def _sse(*events: str) -> bytes:
    return "".join(f"data: {e}\n\n" for e in events).encode()


def _collect(b: BankmlBackend) -> list[str]:
    async def run() -> list[str]:
        out: list[str] = []
        async for tok in await b.stream_chat(_req(stream=True)):
            out.append(tok)
        return out

    return asyncio.run(run())


def test_stream_yields_tokens_and_keeps_the_final_receipt() -> None:
    def chunk(text: str) -> str:
        return json.dumps({"choices": [{"index": 0, "delta": {"content": text},
                                        "finish_reason": None}]})

    final = json.dumps({"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
                        "usage": {"completion_tokens": 2, "prompt_tokens": 12}})
    body = _sse(chunk("I am"), chunk(" mindX."), final,
                json.dumps({"bankml_receipt": RECEIPT}), "[DONE]")

    def handler(request: httpx.Request) -> httpx.Response:
        assert json.loads(request.content)["stream"] is True
        return httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    assert _collect(b) == ["I am", " mindX."]
    assert b.last_receipt == RECEIPT


def test_stream_400_raises_refusal() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(400, text="repeat_penalty: not reproduced")

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    with pytest.raises(BankmlRefusal):
        _collect(b)


def test_stream_error_event_raises() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, content=_sse(json.dumps({"error": "context full"})))

    b = BankmlBackend(transport=httpx.MockTransport(handler))
    with pytest.raises(BankmlError, match="context full"):
        _collect(b)


# ---- operator wiring ---------------------------------------------------------


def test_backend_kwargs_for_bankml(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("MINDXTRAIN_BANKML_BASE_URL", "http://127.0.0.1:18100/v1")
    assert operator_app.backend_kwargs("bankml") == {"base_url": "http://127.0.0.1:18100/v1"}


def test_backend_kwargs_openai_strict_contract(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("MINDXTRAIN_OPENAI_BASE_URL", raising=False)
    with pytest.raises(KeyError):
        operator_app.backend_kwargs("openai_compat", strict=True)
    assert operator_app.backend_kwargs("openai_compat")["base_url"] == ""


def test_bankml_probe_hits_the_identity_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
    seen: list[str] = []

    class _Client:
        def __init__(self, *a: object, **k: object) -> None: ...
        def __enter__(self) -> _Client:
            return self
        def __exit__(self, *a: object) -> None: ...
        def get(self, url: str) -> httpx.Response:
            seen.append(url)
            return httpx.Response(200, json={"bankml": "0.3.4"})

    monkeypatch.delenv("MINDXTRAIN_BANKML_BASE_URL", raising=False)
    monkeypatch.setattr(operator_app.httpx, "Client", _Client)
    assert _REAL_BANKML_REACHABLE() is True
    assert seen == ["http://127.0.0.1:18093/bankml"]


def test_bankml_probe_false_on_connection_error(monkeypatch: pytest.MonkeyPatch) -> None:
    class _Client:
        def __init__(self, *a: object, **k: object) -> None: ...
        def __enter__(self) -> _Client:
            return self
        def __exit__(self, *a: object) -> None: ...
        def get(self, url: str) -> httpx.Response:
            raise httpx.ConnectError("refused")

    monkeypatch.setattr(operator_app.httpx, "Client", _Client)
    assert _REAL_BANKML_REACHABLE() is False


def test_autodetect_order(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("MINDXTRAIN_BACKEND", raising=False)
    monkeypatch.delenv("AUTOMINDX_BACKEND", raising=False)
    monkeypatch.setattr(operator_app, "_bankml_reachable", lambda: True)
    # ollama first, always
    monkeypatch.setattr(operator_app, "_ollama_reachable", lambda: True)
    monkeypatch.setattr(operator_app, "_vllm_reachable", lambda: False)
    assert operator_app.resolve_backend_name() == "ollama"
    # vllm before bankml
    monkeypatch.setattr(operator_app, "_ollama_reachable", lambda: False)
    monkeypatch.setattr(operator_app, "_vllm_reachable", lambda: True)
    assert operator_app.resolve_backend_name() == "vllm"
    # bankml only when it is the one answering
    monkeypatch.setattr(operator_app, "_vllm_reachable", lambda: False)
    assert operator_app.resolve_backend_name() == "bankml"
    # nothing answering: unchanged fallback
    monkeypatch.setattr(operator_app, "_bankml_reachable", lambda: False)
    assert operator_app.resolve_backend_name() == "vllm"


def test_health_dispatch(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(operator_app, "_bankml_reachable", lambda: True)
    monkeypatch.setattr(operator_app, "_bankml_first_model", lambda: "mindx-gen39-F16.gguf")
    assert operator_app.backend_reachable("bankml") is True
    assert operator_app.backend_first_model("bankml") == "mindx-gen39-F16.gguf"


def test_operator_route_passes_refusal_through_as_400(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("MINDXTRAIN_BACKEND", "bankml")

    async def refuse(self: BankmlBackend, request: ChatRequest) -> object:
        raise BankmlRefusal("repeat_penalty: not reproduced")

    monkeypatch.setattr(BankmlBackend, "chat", refuse)
    client = TestClient(operator_app.app)
    r = client.post("/v1/chat/completions", json={
        "model": "mindx-gen39", "messages": [{"role": "user", "content": "hi"}]})
    assert r.status_code == 400
    assert "repeat_penalty: not reproduced" in r.json()["detail"]


def test_operator_route_returns_receipt(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("MINDXTRAIN_BACKEND", "bankml")
    real_client = httpx.AsyncClient
    transport = httpx.MockTransport(lambda req: httpx.Response(200, json=_completion()))
    monkeypatch.setattr(
        "mindxtrain.operator.backends.bankml.httpx.AsyncClient",
        lambda **kw: real_client(timeout=kw.get("timeout"), transport=transport),
    )
    client = TestClient(operator_app.app)
    r = client.post("/v1/chat/completions", json={
        "model": "mindx-gen39", "messages": [{"role": "user", "content": "hi"}]})
    assert r.status_code == 200
    assert r.json()["receipt"]["model_sha256"] == RECEIPT["model_sha256"]


# ---- governance panel env chain -----------------------------------------------


def test_panel_bankml_url_only_when_set_and_last(monkeypatch: pytest.MonkeyPatch) -> None:
    for e in ("MINDXTRAIN_OPENAI_BASE_URL", "MINDXTRAIN_VLLM_BASE_URL",
              "MINDXTRAIN_OLLAMA_BASE_URL", "MINDXTRAIN_BANKML_BASE_URL", "MINDXTRAIN_BACKEND"):
        monkeypatch.delenv(e, raising=False)
    assert P.resolve_chat_base_url() == "http://localhost:11434/v1"
    monkeypatch.setenv("MINDXTRAIN_BANKML_BASE_URL", "http://127.0.0.1:18093/v1")
    assert P.resolve_chat_base_url() == "http://127.0.0.1:18093/v1"
    monkeypatch.setenv("MINDXTRAIN_OLLAMA_BASE_URL", "http://localhost:11434/v1")
    assert P.resolve_chat_base_url() == "http://localhost:11434/v1"
    monkeypatch.setenv("MINDXTRAIN_BACKEND", "bankml")
    assert P.resolve_chat_base_url() == "http://127.0.0.1:18093/v1"


# ---- console + HF Modelfile: neutral defaults are not sent ----------------------


def test_console_wire_options_omits_engine_defaults() -> None:
    from mindxtrain.ui import console

    sent = console.wire_options(dict(console.DEFAULTS))
    for k in ("repeat_penalty", "presence_penalty", "frequency_penalty", "mirostat",
              "mirostat_tau", "mirostat_eta", "stop"):
        assert k not in sent, k
    assert sent["temperature"] == 0.7 and sent["top_k"] == 40
    # a deliberate value is always sent — the engine then honours or refuses it, visibly
    deliberate = console.wire_options({"repeat_penalty": 1.3, "mirostat": 2, "mirostat_tau": 4.0})
    assert deliberate == {"repeat_penalty": 1.3, "mirostat": 2, "mirostat_tau": 4.0}


def test_hf_published_modelfile_repeat_penalty_is_optional() -> None:
    from mindxtrain.hf.extension import published_modelfile

    default = published_modelfile("You are mindX.")
    assert "PARAMETER repeat_penalty 1.3" in default  # default unchanged
    assert default.endswith('PARAMETER stop "<|im_end|>"\n')
    assert "repeat_penalty" not in published_modelfile("You are mindX.", repeat_penalty=None)