evalrx 0.1.2__py3-none-any.whl

evalrx/models/backends/__init__.py

@@ -0,0 +1,58 @@
+"""Backends — runtimes that turn a ModelSpec into a Model.
+
+Capabilities come from the backend; identity from the spec.  All three import
+torch-free (heavy deps are lazy inside ``build``/``load``), so the registry and
+``compose()`` work on the light install.
+"""
+
+from evalrx.models.backends.api import (
+    APIBackend,
+    APIModel,
+    call_vision_api_chat_fn,
+    call_vision_api_generate_fn,
+    parse_openai_logprobs,
+)
+from evalrx.models.backends.base import Backend, RuntimeConfig
+from evalrx.models.backends.gemini_compat import (
+    ThinkingPolicy,
+    gemini_chat_fn,
+    gemini_generate_fn,
+    gemini_logprobs_fn,
+    gemini_runtime,
+)
+from evalrx.models.backends.hf_local import HFLocalBackend, HFLocalModel
+from evalrx.models.backends.openai_compat import (
+    openai_chat_fn,
+    openai_generate_fn,
+    openai_runtime,
+)
+from evalrx.models.backends.vllm_offline import VLLMOfflineBackend
+
+#: name -> backend class.  Plain dict (no import-side-effect decorator).
+BACKENDS: dict[str, type[Backend]] = {
+    "api": APIBackend,
+    "hf_local": HFLocalBackend,
+    "vllm_offline": VLLMOfflineBackend,
+}
+
+__all__ = [
+    "Backend",
+    "RuntimeConfig",
+    "APIBackend",
+    "APIModel",
+    "HFLocalBackend",
+    "HFLocalModel",
+    "VLLMOfflineBackend",
+    "call_vision_api_generate_fn",
+    "call_vision_api_chat_fn",
+    "parse_openai_logprobs",
+    "openai_chat_fn",
+    "openai_generate_fn",
+    "openai_runtime",
+    "ThinkingPolicy",
+    "gemini_chat_fn",
+    "gemini_generate_fn",
+    "gemini_logprobs_fn",
+    "gemini_runtime",
+    "BACKENDS",
+]

evalrx/models/backends/api.py

@@ -0,0 +1,169 @@
+"""API backend — remote / OpenAI-compatible models (incl. ``vllm serve``).
+
+Capabilities: ``GENERATE`` + ``TOOL_CALLS`` (+ ``LOGPROBS`` when the endpoint
+returns top-logprobs, e.g. OpenAI but not Gemini).  No internals: ``forward``
+raises, and capability matching means white-box analyzers are simply never
+offered for an API model.
+
+Reuse of your existing engine is by **dependency injection**: pass your
+``call_vision_api`` adapted through :func:`call_vision_api_generate_fn` as
+``RuntimeConfig.generate_fn`` — no cross-repo import, no rewrite.  A ``vllm serve``
+endpoint is just this backend pointed at ``localhost`` (the answer to "is
+vllm_local the same as api?": *for serving, yes* — see ``vllm_offline`` for the
+distinct in-process batch path).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from evalrx.core.capability import Capability, CapabilityError
+from evalrx.core.case import Inputs
+from evalrx.core.model import Model, TokenLogprob, Trace
+from evalrx.core.tool import ChatTurn
+from evalrx.models.backends.base import Backend, RuntimeConfig
+
+
+def parse_openai_logprobs(content_logprobs) -> list[TokenLogprob]:
+    """Parse OpenAI ``choices[0].logprobs.content`` into ``list[TokenLogprob]``.
+
+    Use inside a ``logprobs_fn`` after calling an OpenAI-compatible endpoint with
+    ``logprobs=True, top_logprobs=k``.
+    """
+    out: list[TokenLogprob] = []
+    for item in content_logprobs or []:
+        top = {t["token"]: float(t["logprob"]) for t in item.get("top_logprobs", [])}
+        out.append(TokenLogprob(token=item.get("token", ""), logprob=float(item["logprob"]), top=top))
+    return out
+
+_WHITEBOX = {
+    Capability.ATTENTION,
+    Capability.HIDDEN_STATES,
+    Capability.LOGITS,
+    Capability.GRADIENTS,
+    Capability.ACTIVATIONS,
+    Capability.EMBEDDINGS,
+}
+
+
+def call_vision_api_generate_fn(
+    call_vision_api: Callable,
+    *,
+    default_sampling: dict | None = None,
+) -> Callable[..., str]:
+    """Adapt an existing ``call_vision_api(model, messages, sampling_params, tools)``
+    into the simple ``generate_fn(prompt, model=..., **kw) -> str`` an APIModel needs.
+
+    This is the one-liner that reuses your XSkill engine verbatim::
+
+        from engine.api_caller import call_vision_api
+        rt = RuntimeConfig(generate_fn=call_vision_api_generate_fn(call_vision_api))
+    """
+
+    def _fn(prompt: str, model: str, **kw) -> str:
+        messages = [{"role": "user", "content": prompt}]
+        sampling = {**(default_sampling or {}), **kw}
+        resp = call_vision_api(model, messages, sampling)
+        if isinstance(resp, dict):  # tool/reasoning shape -> pull text
+            return resp.get("content") or str(resp)
+        return resp
+
+    return _fn
+
+
+def call_vision_api_chat_fn(call_vision_api: Callable) -> Callable[..., ChatTurn]:
+    """Adapt ``call_vision_api(model, messages, sampling, tools)`` into a tool-aware
+    ``chat_fn(messages, tools, model) -> ChatTurn`` for agent mode.
+
+    The engine already returns structured tool_calls when present, so this just
+    normalises its reply shape::
+
+        rt = RuntimeConfig(chat_fn=call_vision_api_chat_fn(call_vision_api))
+    """
+
+    def _fn(messages: list, tools=None, model: str = "") -> ChatTurn:
+        resp = call_vision_api(model, messages, {}, tools=tools)
+        if isinstance(resp, dict):
+            return ChatTurn(text=resp.get("content") or "", raw_tool_calls=resp.get("tool_calls"))
+        return ChatTurn(text=resp, raw_tool_calls=None)
+
+    return _fn
+
+
+class APIModel(Model):
+    """A model reachable only through a text-in / text-out API."""
+
+    def __init__(self, spec, runtime: RuntimeConfig) -> None:
+        self.spec = spec
+        self.runtime = runtime
+        self._generate_fn = runtime.generate_fn
+        self._chat_fn = runtime.chat_fn
+        self._logprobs_fn = runtime.logprobs_fn
+        caps = {Capability.GENERATE, Capability.TOOL_CALLS}
+        if runtime.logprobs_fn is not None:  # only claim LOGPROBS if we can actually retrieve them
+            caps.add(Capability.LOGPROBS)
+        self.capabilities = frozenset(caps)
+        self.modalities = frozenset(spec.modalities)  # text / image / audio / video, from the spec
+
+    def generate(self, inputs: Any, **kwargs) -> str:
+        if self._generate_fn is None:
+            raise RuntimeError(
+                f"APIModel({self.spec.key!r}) has no generate_fn. Inject one via "
+                "RuntimeConfig(generate_fn=...), e.g. call_vision_api_generate_fn(call_vision_api)."
+            )
+        prompt = inputs.prompt if isinstance(inputs, Inputs) else str(inputs)
+        model_name = self.spec.hf_repo or self.spec.key
+        return self._generate_fn(prompt, model=model_name, **self._media(inputs), **kwargs)
+
+    def chat(self, messages: list, tools=None) -> ChatTurn:
+        model_name = self.spec.hf_repo or self.spec.key
+        if self._chat_fn is not None:
+            return self._chat_fn(messages, tools=tools, model=model_name)
+        # Fallback: no tool-aware chat wired -> text-only (no tool calls possible).
+        if self._generate_fn is None:
+            raise RuntimeError(
+                f"APIModel({self.spec.key!r}) needs a chat_fn (tool-aware) or generate_fn. "
+                "Inject one, e.g. call_vision_api_chat_fn(call_vision_api)."
+            )
+        prompt = messages[-1].get("content", "") if messages else ""
+        return ChatTurn(text=self._generate_fn(prompt, model=model_name), raw_tool_calls=None)
+
+    def logprobs(self, inputs: Any, **kwargs) -> list[TokenLogprob]:
+        if self._logprobs_fn is None:
+            raise CapabilityError(analyzer="logprobs", model=repr(self), missing={Capability.LOGPROBS})
+        prompt = inputs.prompt if isinstance(inputs, Inputs) else str(inputs)
+        return self._logprobs_fn(
+            prompt, model=self.spec.hf_repo or self.spec.key, **self._media(inputs), **kwargs
+        )
+
+    def _media(self, inputs: Any) -> dict:
+        """The filled media slots of *inputs* that the spec declares.
+
+        ``image`` / ``audio`` go to ``generate_fn`` / ``logprobs_fn`` as
+        keyword arguments; a text-only spec never forwards them, and a slot
+        the case leaves empty is not sent. Without this the api backend
+        answered every VLM/ALM case from the prompt alone.
+        """
+        if not isinstance(inputs, Inputs):
+            return {}
+        media = {}
+        if "image" in self.modalities and inputs.image is not None:
+            media["image"] = inputs.image
+        if "audio" in self.modalities and inputs.audio is not None:
+            media["audio"] = inputs.audio
+        return media
+
+    def forward(self, inputs: Any, capture: set[Capability], spec=None) -> Trace:
+        missing = set(capture) & _WHITEBOX
+        raise CapabilityError(analyzer="forward", model=repr(self), missing=missing or set(capture))
+
+    def __repr__(self) -> str:
+        return f"APIModel(key={self.spec.key!r}, caps={sorted(c.value for c in self.capabilities)})"
+
+
+class APIBackend(Backend):
+    kind = "api"
+    capabilities = frozenset({Capability.GENERATE, Capability.TOOL_CALLS, Capability.LOGPROBS})
+
+    def build(self, spec, runtime: RuntimeConfig) -> APIModel:
+        return APIModel(spec, runtime)

evalrx/models/backends/base.py

@@ -0,0 +1,57 @@
+"""Backend — a runtime that turns a ModelSpec into a Model.
+
+The orthogonal decomposition: **capabilities come from the backend**,
+**identity comes from the spec**.  The same :class:`~evalrx.core.spec.ModelSpec`
+drives an API call (``api``), a high-throughput logprob batch (``vllm_offline``),
+or a full internals capture (``hf_local``) — only the capability set differs.
+
+This module is torch-free; concrete backends import their heavy deps lazily.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+from evalrx.core.capability import Capability
+
+if TYPE_CHECKING:
+    from evalrx.core.model import Model
+    from evalrx.core.spec import ModelSpec
+
+
+@dataclass
+class RuntimeConfig:
+    """How to run a model — orthogonal to *which* model (that's the spec)."""
+
+    device: str = "auto"
+    dtype: str = "bfloat16"
+    attn_impl: Optional[str] = None          # hf_local forces "eager" when attention is wanted
+    max_new_tokens: int = 512
+    #: hf_local, TEXT-ONLY specs: render ``inputs.prompt`` as one user turn through
+    #: the tokenizer's chat template (with ``spec.chat_template_kwargs``, e.g.
+    #: ``enable_thinking=False``) before generate()/logprobs(). Off by default: a
+    #: raw prompt is tokenised verbatim (completion mode), which is what every
+    #: caller got before 2026-08-21 and what a pre-formatted prompt needs. The
+    #: multimodal encode path always renders the template.
+    apply_chat_template: bool = False
+    engine_kwargs: dict[str, Any] = field(default_factory=dict)   # vllm_offline LLM(**)
+    client_kwargs: dict[str, Any] = field(default_factory=dict)   # api client opts (e.g. logprobs=True)
+    generate_fn: Optional[Callable[..., str]] = None              # api: simple text generate
+    chat_fn: Optional[Callable[..., Any]] = None                  # api: tool-aware chat (returns ChatTurn)
+    logprobs_fn: Optional[Callable[..., Any]] = None              # api: returns list[TokenLogprob]
+
+
+class Backend(ABC):
+    """Base for runtimes.  Declares the capability set it can provide."""
+
+    kind: str = "base"
+    capabilities: frozenset[Capability] = frozenset()
+
+    @abstractmethod
+    def build(self, spec: "ModelSpec", runtime: RuntimeConfig) -> "Model":
+        """Construct a :class:`~evalrx.core.model.Model` for *spec* (lazy load)."""
+
+    def __repr__(self) -> str:  # pragma: no cover - cosmetic
+        return f"{type(self).__name__}(kind={self.kind!r})"