Blame for mindxtrain/operator/backends/bankml.py

"""bankml backend — the verified CPU engine, reached over HTTP only.

bankml (https://github.com/cryptoAGI/bankml, Apache-2.0 OR MIT) is a zero-dependency Rust
runtime whose `bankml serve --native` answers OpenAI `/v1/chat/completions` and Ollama `/api/*`
from its own forward pass, token-identical to llama.cpp b11192, on 127.0.0.1:18093 by default.
mindXtrain never vendors it: this module speaks its documented wire protocol and nothing else
(clean-room policy, CLAUDE.md).

What this backend adds over `openai_compat`:

- **Receipts.** Every bankml answer carries a `bankml_receipt` (model_sha256, request_sha256,
  response_sha256, engine, tokens, wall_ms, ...). Non-streamed it is a top-level object and lands
  in `ChatResponse.receipt`; streamed it arrives as one extra `data: {"bankml_receipt": ...}`
  event before `data: [DONE]`. Either way the latest receipt is also kept on `last_receipt`.
- **Refusals are typed, never retried.** bankml answers HTTP 400 with a plain-text reason when a
  request asks for something its verified forward pass does not reproduce (penalties before
  0.3.6, mirostat, typical_p, tools, an unknown architecture, Q8_0 / Q4_K / BF16).
  That becomes `BankmlRefusal(reason)`. The request is never re-sent with altered parameters:
  a changed request would be a different, unreceipted question.

- **Sampling options, only when set.** bankml 0.3.6 reproduces `repeat_penalty`,
  `repeat_last_n`, `presence_penalty` and `frequency_penalty` as llama-server b11192 does; a
  small generation (mindx-genN) needs `repeat_penalty` to answer in words. They are sent from
  `options=` or `MINDXTRAIN_BANKML_OPTIONS`, never by default, so an older bankml is never sent
  a field it refuses.

Env: `MINDXTRAIN_BANKML_BASE_URL` (default `http://127.0.0.1:18093/v1`),
`MINDXTRAIN_BANKML_OPTIONS` (a JSON object, e.g. `{"repeat_penalty": 1.3}`).
"""

from __future__ import annotations

import json
import os
from collections.abc import AsyncIterator
from typing import Any

import httpx

from mindxtrain.models.registry import ChatRequest, ChatResponse, register_backend
from mindxtrain.operator.backends.openai_compat import OpenAICompatBackend

DEFAULT_BANKML_BASE_URL = "http://127.0.0.1:18093/v1"


def bankml_base_url() -> str:
    """The configured bankml OpenAI base URL (`.../v1`)."""
    return os.environ.get("MINDXTRAIN_BANKML_BASE_URL", DEFAULT_BANKML_BASE_URL).rstrip("/")


def bankml_root_url(base_url: str | None = None) -> str:
    """The server root (no `/v1`) — where `/health`, `/bankml` and `/api/*` live."""
    return (base_url or bankml_base_url()).rstrip("/").removesuffix("/v1")


class BankmlError(RuntimeError):
    """bankml answered, but not with an answer (an engine error, a broken stream)."""

    label = "error"

    def __init__(self, reason: str, status_code: int = 0) -> None:
        self.reason = reason.strip()
        self.status_code = status_code
        super().__init__(f"bankml {self.label} ({status_code}): {self.reason}")


class BankmlRefusal(BankmlError):
    """bankml refused the request (HTTP 400) and said why.

    The reason is bankml's own text, e.g. ``mirostat: not reproduced; ...``. Callers must not
    retry with parameters stripped or altered — change the request deliberately instead.
    """

    label = "refused"

    def __init__(self, reason: str, status_code: int = 400) -> None:
        super().__init__(reason, status_code)


def refusal_reason(resp: httpx.Response) -> str:
    """bankml's reason from an error body: `/v1` sends text/plain, `/api` sends `{"error": ...}`."""
    text = resp.text or ""
    try:
        body = json.loads(text)
    except ValueError:
        return text.strip() or f"HTTP {resp.status_code}"
    if isinstance(body, dict) and body.get("error"):
        err = body["error"]
        return str(err.get("message", err) if isinstance(err, dict) else err)
    return text.strip() or f"HTTP {resp.status_code}"


def raise_for_bankml(resp: httpx.Response) -> None:
    """400 → `BankmlRefusal`; any other non-2xx → `BankmlError`."""
    if resp.status_code == 400:
        raise BankmlRefusal(refusal_reason(resp))
    if resp.status_code >= 400:
        raise BankmlError(refusal_reason(resp), resp.status_code)


def _env_options() -> dict[str, object]:
    raw = os.environ.get("MINDXTRAIN_BANKML_OPTIONS")
    if not raw:
        return {}
    val = json.loads(raw)
    if not isinstance(val, dict):
        msg = "MINDXTRAIN_BANKML_OPTIONS must be a JSON object"
        raise ValueError(msg)
    return val


@register_backend("bankml")
class BankmlBackend(OpenAICompatBackend):
    """OpenAI-compatible client for `bankml serve --native`, keeping receipts and refusals."""

    name = "bankml"

    def __init__(
        self,
        base_url: str | None = None,
        timeout_s: float = 300.0,
        *,
        seed: int | None = None,
        options: dict[str, object] | None = None,
        transport: httpx.AsyncBaseTransport | None = None,
    ) -> None:
        super().__init__(base_url=base_url or bankml_base_url(), api_key=None, timeout_s=timeout_s)
        # `api_key or env` in the parent would pick up an OpenAI key; bankml is loopback, keyless.
        self.api_key = ""
        self.seed = seed
        # Sampling fields beyond OpenAI's that bankml honours (repeat_penalty & co. from 0.3.6),
        # sent only when set: from MINDXTRAIN_BANKML_OPTIONS (a JSON object), then `options`.
        self.options: dict[str, object] = {**_env_options(), **(options or {})}
        self._transport = transport
        self.last_receipt: dict[str, Any] | None = None

    def _client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(timeout=self.timeout_s, transport=self._transport)

    def _payload(self, request: ChatRequest, *, stream: bool) -> dict[str, object]:
        payload = {**self.options, **super()._payload(request, stream=stream)}
        if self.seed is not None:
            payload["seed"] = self.seed
        return payload

    async def chat(self, request: ChatRequest) -> ChatResponse:
        async with self._client() as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                json=self._payload(request, stream=False),
                headers=self._headers(),
            )
        raise_for_bankml(resp)
        data = resp.json()
        choice = (data.get("choices") or [{}])[0]
        usage = data.get("usage") or {}
        receipt = data.get("bankml_receipt")
        self.last_receipt = receipt if isinstance(receipt, dict) else None
        finish = choice.get("finish_reason") or "stop"
        return ChatResponse(
            model=data.get("model", request.model),
            content=(choice.get("message") or {}).get("content", "") or "",
            finish_reason=finish if finish in ("stop", "length", "error") else "stop",
            prompt_tokens=int(usage.get("prompt_tokens", 0)),
            completion_tokens=int(usage.get("completion_tokens", 0)),
            receipt=self.last_receipt,
        )

    async def stream_chat(self, request: ChatRequest) -> AsyncIterator[str]:
        self.last_receipt = None

        async def _gen() -> AsyncIterator[str]:
            async with self._client() as client:
                async with client.stream(
                    "POST",
                    f"{self.base_url}/chat/completions",
                    json=self._payload(request, stream=True),
                    headers=self._headers(),
                ) as resp:
                    if resp.status_code >= 400:
                        await resp.aread()
                        raise_for_bankml(resp)
                    async for raw in resp.aiter_lines():
                        if not raw or not raw.startswith("data:"):
                            continue
                        data = raw[5:].strip()
                        if data == "[DONE]":
                            return
                        try:
                            chunk = json.loads(data)
                        except json.JSONDecodeError:
                            continue
                        if not isinstance(chunk, dict):
                            continue
                        if "bankml_receipt" in chunk:
                            rec = chunk["bankml_receipt"]
                            self.last_receipt = rec if isinstance(rec, dict) else None
                            continue
                        if "error" in chunk:
                            # bankml stops a stream it cannot finish with `data: {"error": ...}`
                            raise BankmlError(str(chunk["error"]), resp.status_code)
                        delta = (chunk.get("choices") or [{}])[0].get("delta", {})
                        token = delta.get("content")
                        if token:
                            yield token

        return _gen()


__all__ = [
    "DEFAULT_BANKML_BASE_URL",
    "BankmlBackend",
    "BankmlError",
    "BankmlRefusal",
    "bankml_base_url",
    "bankml_root_url",
    "raise_for_bankml",
    "refusal_reason",
]