bankml: the verified CPU engine as backend, serve target and imprint probe (#1)

730c5bb about 1 month ago

9.13 kB

"""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",
]

	"""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",
	]