jev-compatible-server 0.1.0__py3-none-any.whl

jev_compatible_server/registry.py

@@ -0,0 +1,283 @@
+"""Configurable model registry for models with stale or incomplete metadata."""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from importlib import resources
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .backends import (
+    LlamaBackend,
+    PointerTransformersBackend,
+    TransformersBackend,
+    load_decision_config,
+)
+from .bosun import BosunDecisionBackend
+from .causal_options import CausalOptionsBackend
+from .classifier_adapters import GLiClassCalibratedBackend, NLIEntailmentBackend
+from .cross_encoder import CrossEncoderBackend
+from .custom_heads import OpenJevScalarHeadBackend, SmallJevSemanticBackend
+from .encoder_decoder import EncoderDecoderMarginBackend, decision_metadata
+from .gliner2 import GLiNER2Runtime
+from .hidden_state_probe import HiddenStateProbeBackend
+from .laya import LayaBackend
+from .native_systemone import (
+    DjevHTTPRuntime,
+    DjevThinkingRuntime,
+    JeffHTTPRuntime,
+    OpenJevThinkingHTTPRuntime,
+    WinnowHTTPRuntime,
+)
+from .protocol import DecisionRequest, DecisionResponse, UnsupportedAnswer, Usage
+from .runtime import (
+    DecisionRuntime,
+    RuntimeErrorBase,
+    apply_question_type_support,
+    configured_question_types,
+)
+from .sequence_classifier import SequenceClassifierMarginBackend
+
+
+class RegistryModel(BaseModel):
+    """One deployable model and its service-owned decision metadata."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    backend: Literal["llama", "transformers", "mlx"]
+    model: str
+    recipe: str | None = None
+    config: dict[str, Any] = Field(default_factory=dict)
+    config_path: str | None = None
+    enabled: bool = True
+    description: str | None = None
+    support_status: Literal["supported", "pending"] = "supported"
+
+    def resolved_config(
+        self, recipe_config: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        resolved = _merge_config({}, recipe_config or {})
+        file_config = load_decision_config(self.config_path) if self.config_path else {}
+        # Registry values win over model-published metadata and file defaults.
+        resolved = _merge_config(resolved, file_config)
+        return _merge_config(resolved, self.config)
+
+
+class RegistryFile(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    default: str | None = None
+    recipes: dict[str, dict[str, Any]] = Field(default_factory=dict)
+    models: dict[str, RegistryModel] = Field(min_length=1)
+
+
+def _merge_config(base: dict[str, Any], override: Mapping[str, Any]) -> dict[str, Any]:
+    result = dict(base)
+    for key, value in override.items():
+        existing = result.get(key)
+        if isinstance(existing, dict) and isinstance(value, dict):
+            result[key] = _merge_config(existing, value)
+        else:
+            result[key] = value
+    return result
+
+
+class ModelRegistry:
+    def __init__(self, definition: RegistryFile):
+        self.definition = definition
+
+    @classmethod
+    def from_file(cls, path: str | Path) -> ModelRegistry:
+        return cls.from_json(Path(path).read_text())
+
+    @classmethod
+    def from_builtin(cls) -> ModelRegistry:
+        resource = resources.files("jev_compatible_server").joinpath("public-models.json")
+        try:
+            value = resource.read_text()
+        except FileNotFoundError:
+            source_path = Path(__file__).parents[2] / "configs" / "public-models.json"
+            value = source_path.read_text()
+        return cls.from_json(value)
+
+    @classmethod
+    def from_json(cls, value: str) -> ModelRegistry:
+        raw = json.loads(value)
+        if not isinstance(raw, dict):
+            raise RuntimeErrorBase("model registry must be a JSON object")
+        try:
+            definition = RegistryFile.model_validate(raw)
+        except ValueError as exc:
+            raise RuntimeErrorBase(f"invalid model registry: {exc}") from exc
+        if definition.default is not None and definition.default not in definition.models:
+            raise RuntimeErrorBase(f"registry default is not registered: {definition.default}")
+        for name, entry in definition.models.items():
+            if entry.recipe is not None and entry.recipe not in definition.recipes:
+                raise RuntimeErrorBase(
+                    f"model {name!r} references unregistered recipe: {entry.recipe}"
+                )
+        return cls(definition)
+
+    def resolve(self, requested: str | None) -> tuple[str, RegistryModel]:
+        name = requested or self.definition.default
+        if name is None:
+            raise RuntimeErrorBase("request must specify model or registry needs a default")
+        entry = self.definition.models.get(name)
+        if entry is None:
+            raise RuntimeErrorBase(f"model is not registered: {name}")
+        if not entry.enabled:
+            raise RuntimeErrorBase(f"model is disabled: {name}")
+        return name, entry
+
+
+def build_transformers_runtime(
+    model_id: str,
+    config: dict[str, Any],
+    *,
+    batch_size_override: int | None = None,
+) -> DecisionRuntime:
+    """Select a Transformers readout from metadata, never from a model name."""
+
+    effective_config = dict(config)
+    if batch_size_override is not None:
+        if isinstance(batch_size_override, bool) or batch_size_override <= 0:
+            raise RuntimeErrorBase("model batch size override must be a positive integer")
+        # Dotted metadata has highest precedence in decision_metadata(), so this
+        # overrides both nested and dotted registry recipe values without
+        # mutating the registry's reusable configuration.
+        effective_config["decision.batch_size"] = batch_size_override
+
+    readout = decision_metadata(effective_config).get("readout")
+    if readout == "pointer_head":
+        return PointerTransformersBackend(model_id, config=effective_config)
+    if readout == "encoder_decoder_margin":
+        return EncoderDecoderMarginBackend(model_id, config=effective_config)
+    if readout == "sequence_classifier_margin":
+        return SequenceClassifierMarginBackend(model_id, config=effective_config)
+    if readout == "cross_encoder_margin":
+        return CrossEncoderBackend(model_id, config=effective_config)
+    if readout == "causal_options":
+        return CausalOptionsBackend(model_id, config=effective_config)
+    if readout == "bosun_decision_tokens":
+        return BosunDecisionBackend(model_id, config=effective_config)
+    if readout == "nli_entailment":
+        return NLIEntailmentBackend(model_id, config=effective_config)
+    if readout == "gliclass_calibrated":
+        return GLiClassCalibratedBackend(model_id, config=effective_config)
+    if readout == "openjev_scalar_head":
+        return OpenJevScalarHeadBackend(model_id, config=effective_config)
+    if readout == "semantic_option_head":
+        return SmallJevSemanticBackend(model_id, config=effective_config)
+    if readout == "hidden_state_probe":
+        return HiddenStateProbeBackend(model_id, config=effective_config)
+    if readout == "laya_native":
+        return LayaBackend(model_id, config=effective_config)
+    if readout == "diffusion_structured_read":
+        return DjevHTTPRuntime(model_id, config=effective_config)
+    if readout == "diffusion_thinking_read":
+        native_contract = decision_metadata(effective_config).get("native_contract")
+        if native_contract == "djev":
+            return DjevThinkingRuntime(model_id, config=effective_config)
+        if native_contract == "openjev":
+            return OpenJevThinkingHTTPRuntime(model_id, config=effective_config)
+        raise RuntimeErrorBase(
+            "diffusion_thinking_read requires decision.native_contract of 'djev' or 'openjev'"
+        )
+    if readout == "gliformer_native":
+        return JeffHTTPRuntime(model_id, config=effective_config)
+    if readout == "gliner2_multilabel":
+        return GLiNER2Runtime(model_id, config=effective_config)
+    return TransformersBackend(model_id, config=effective_config)
+
+
+class RegistryRuntime(DecisionRuntime):
+    """Lazy, cached runtime dispatch with per-model microbatching."""
+
+    def __init__(
+        self,
+        registry: ModelRegistry,
+        *,
+        model_batch_size: int | None = None,
+    ):
+        self.registry = registry
+        self.model_name = registry.definition.default or "registry"
+        self.model_batch_size = model_batch_size
+        self._runtimes: dict[str, DecisionRuntime] = {}
+
+    def _runtime(self, name: str, entry: RegistryModel) -> DecisionRuntime:
+        cached = self._runtimes.get(name)
+        if cached is not None:
+            return cached
+        recipe_config = (
+            self.registry.definition.recipes.get(entry.recipe, {})
+            if entry.recipe is not None
+            else {}
+        )
+        config = entry.resolved_config(recipe_config)
+        runtime: DecisionRuntime
+        if entry.backend == "llama":
+            if decision_metadata(config).get("readout") == "winnow_shared_branch":
+                runtime = WinnowHTTPRuntime(entry.model, config=config)
+            else:
+                runtime = LlamaBackend(entry.model, config=config)
+        elif entry.backend == "mlx":
+            raise RuntimeErrorBase(
+                "the MLX backend is registered but not installed in this service image"
+            )
+        else:
+            runtime = build_transformers_runtime(
+                entry.model,
+                config,
+                batch_size_override=self.model_batch_size,
+            )
+        runtime = apply_question_type_support(runtime, config)
+        self._runtimes[name] = runtime
+        return runtime
+
+    def decide_batch(self, requests: Sequence[DecisionRequest]) -> list[DecisionResponse]:
+        grouped: dict[str, list[tuple[int, DecisionRequest]]] = defaultdict(list)
+        entries: dict[str, RegistryModel] = {}
+        for index, request in enumerate(requests):
+            name, entry = self.registry.resolve(request.model)
+            grouped[name].append((index, request))
+            entries[name] = entry
+
+        results: list[DecisionResponse | None] = [None] * len(requests)
+        for name, indexed_requests in grouped.items():
+            entry = entries[name]
+            recipe_config = (
+                self.registry.definition.recipes.get(entry.recipe, {})
+                if entry.recipe is not None
+                else {}
+            )
+            config = entry.resolved_config(recipe_config)
+            supported_types = configured_question_types(config)
+            supported_set = frozenset(supported_types)
+            if not any(
+                question.type in supported_set
+                for _, request in indexed_requests
+                for question in request.questions.values()
+            ):
+                for index, request in indexed_requests:
+                    results[index] = DecisionResponse(
+                        model=entry.model,
+                        answers={
+                            question_name: UnsupportedAnswer(
+                                type="unsupported",
+                                question_type=question.type,
+                                supported_types=list(supported_types),
+                            )
+                            for question_name, question in request.questions.items()
+                        },
+                        usage=Usage(),
+                    )
+                continue
+            runtime = self._runtime(name, entries[name])
+            responses = runtime.decide_batch([request for _, request in indexed_requests])
+            for (index, _), response in zip(indexed_requests, responses, strict=True):
+                results[index] = response
+        return [result for result in results if result is not None]

jev_compatible_server/runtime.py

@@ -0,0 +1,237 @@
+"""Runtime interface and shared decision post-processing."""
+
+from __future__ import annotations
+
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any, cast
+
+from .protocol import (
+    ChoiceAnswer,
+    ChoiceQuestion,
+    DecisionRequest,
+    DecisionResponse,
+    NoulAnswer,
+    NoulQuestion,
+    QuestionType,
+    ScoreAnswer,
+    ScoreQuestion,
+    UnsupportedAnswer,
+    Usage,
+)
+
+
+class RuntimeErrorBase(RuntimeError):
+    """A user-facing runtime configuration or inference error."""
+
+
+def softmax(values: Sequence[float]) -> list[float]:
+    if not values:
+        raise RuntimeErrorBase("cannot normalize an empty score vector")
+    peak = max(values)
+    exps = [math.exp(value - peak) for value in values]
+    total = sum(exps)
+    return [value / total for value in exps]
+
+
+class DecisionRuntime(ABC):
+    """Backend contract. Implementations may process the whole batch at once."""
+
+    model_name: str
+
+    @abstractmethod
+    def decide_batch(
+        self, requests: Sequence[DecisionRequest]
+    ) -> list[DecisionResponse]:
+        """Evaluate requests, preserving input order."""
+
+    def decide(self, request: DecisionRequest) -> DecisionResponse:
+        return self.decide_batch([request])[0]
+
+
+ALL_QUESTION_TYPES: tuple[QuestionType, ...] = ("choice", "score", "noul")
+
+
+def configured_question_types(config: dict[str, Any]) -> tuple[QuestionType, ...]:
+    """Read a model's declared capability boundary from decision metadata."""
+
+    nested = config.get("decision", {})
+    nested_types = nested.get("question_types") if isinstance(nested, dict) else None
+    value = config.get("decision.question_types", nested_types)
+    if value is None:
+        return ALL_QUESTION_TYPES
+    if not isinstance(value, list) or not value:
+        raise RuntimeErrorBase("decision.question_types must be a non-empty array")
+    if any(item not in ALL_QUESTION_TYPES for item in value):
+        raise RuntimeErrorBase(
+            "decision.question_types may contain only choice, score, and noul"
+        )
+    result: list[QuestionType] = []
+    for item in value:
+        typed_item = cast(QuestionType, item)
+        if typed_item not in result:
+            result.append(typed_item)
+    return tuple(result)
+
+
+class QuestionTypeRuntime(DecisionRuntime):
+    """Return explicit per-question unsupported results around any backend."""
+
+    def __init__(
+        self,
+        runtime: DecisionRuntime,
+        supported_types: Sequence[QuestionType],
+    ) -> None:
+        self.runtime = runtime
+        self.model_name = runtime.model_name
+        self.supported_types = tuple(supported_types)
+        self._supported = frozenset(supported_types)
+
+    def decide_batch(
+        self, requests: Sequence[DecisionRequest]
+    ) -> list[DecisionResponse]:
+        filtered: list[DecisionRequest] = []
+        filtered_indices: list[int] = []
+        for index, request in enumerate(requests):
+            questions = {
+                name: question
+                for name, question in request.questions.items()
+                if question.type in self._supported
+            }
+            if questions:
+                filtered.append(request.model_copy(update={"questions": questions}))
+                filtered_indices.append(index)
+
+        inferred: dict[int, DecisionResponse] = {}
+        if filtered:
+            responses = self.runtime.decide_batch(filtered)
+            if len(responses) != len(filtered):
+                raise RuntimeErrorBase("runtime returned the wrong batch length")
+            inferred = dict(zip(filtered_indices, responses, strict=True))
+
+        results: list[DecisionResponse] = []
+        for index, request in enumerate(requests):
+            response = inferred.get(index)
+            supported_answers = response.answers if response is not None else {}
+            answers: dict[str, Any] = {}
+            for name, question in request.questions.items():
+                if question.type in self._supported:
+                    answer = supported_answers.get(name)
+                    if answer is None:
+                        raise RuntimeErrorBase(
+                            f"runtime did not return answer for question {name!r}"
+                        )
+                    answers[name] = answer
+                else:
+                    answers[name] = UnsupportedAnswer(
+                        type="unsupported",
+                        question_type=question.type,
+                        supported_types=list(self.supported_types),
+                    )
+            results.append(
+                DecisionResponse(
+                    model=response.model if response is not None else self.model_name,
+                    answers=answers,
+                    usage=response.usage if response is not None else Usage(),
+                )
+            )
+        return results
+
+
+def apply_question_type_support(
+    runtime: DecisionRuntime, config: dict[str, Any]
+) -> DecisionRuntime:
+    supported = configured_question_types(config)
+    if supported == ALL_QUESTION_TYPES:
+        return runtime
+    return QuestionTypeRuntime(runtime, supported)
+
+
+class TokenLogitRuntime(DecisionRuntime):
+    """Shared formatter for backends that expose next-token logits."""
+
+    def __init__(self, *, model_name: str, config: dict[str, Any] | None = None):
+        self.model_name = model_name
+        self.config = config or {}
+
+    def _metadata(self, name: str, legacy_name: str) -> Any:
+        value = self.config.get(name, self.config.get(legacy_name))
+        if isinstance(value, dict) and name.startswith("decision."):
+            return value
+        return value
+
+    def _prompt(self, request: DecisionRequest, question_name: str, question: Any) -> str:
+        template = self._metadata("decision.prompt_template", "prompt_template")
+        if not isinstance(template, str):
+            raise RuntimeErrorBase(
+                "model metadata must define a string decision prompt_template"
+            )
+        return template.format(
+            state=request.state,
+            question_name=question_name,
+            instructions=question.instructions,
+            criteria=question.criteria if hasattr(question, "criteria") else "",
+        )
+
+    def _token_id(self, label: str) -> int:
+        tokens = self._metadata("decision.tokens", "decision_tokens")
+        if isinstance(tokens, str):
+            import json
+
+            try:
+                tokens = json.loads(tokens)
+            except json.JSONDecodeError as exc:
+                raise RuntimeErrorBase("decision.tokens is not valid JSON") from exc
+        if not isinstance(tokens, dict) or label not in tokens:
+            raise RuntimeErrorBase(f"missing decision token mapping for {label!r}")
+        token_id = tokens[label]
+        if not isinstance(token_id, int):
+            raise RuntimeErrorBase(f"decision token for {label!r} is not an integer")
+        return token_id
+
+    def _answer(self, question: Any, logits: dict[int, float]) -> Any:
+        labels = self._labels_for_question(question)
+        return self.answer_from_label_scores(
+            question, {label: logits[self._token_id(label)] for label in labels}
+        )
+
+    @staticmethod
+    def _labels_for_question(question: Any) -> list[str]:
+        if isinstance(question, ChoiceQuestion):
+            return list(question.criteria)
+        if isinstance(question, ScoreQuestion):
+            return [str(index) for index in range(len(question.criteria))]
+        if isinstance(question, NoulQuestion):
+            return ["true", "false"]
+        raise RuntimeErrorBase(f"unsupported question type: {type(question).__name__}")
+
+    def answer_from_label_scores(self, question: Any, scores: dict[str, float]) -> Any:
+        if isinstance(question, ChoiceQuestion):
+            labels = list(question.criteria)
+            probabilities = softmax([scores[label] for label in labels])
+            distribution = dict(zip(labels, probabilities, strict=True))
+            choice = max(distribution, key=distribution.__getitem__)
+            return ChoiceAnswer(
+                type="choice",
+                choice=choice,
+                probabilities=distribution,
+                confidence=distribution[choice],
+            )
+        if isinstance(question, ScoreQuestion):
+            labels = [str(index) for index in range(len(question.criteria))]
+            probabilities = softmax([scores[label] for label in labels])
+            distribution = dict(zip(labels, probabilities, strict=True))
+            score = sum(index * probability for index, probability in enumerate(probabilities))
+            return ScoreAnswer(
+                type="score",
+                score=score,
+                probabilities=distribution,
+                confidence=max(probabilities),
+                legend=question.criteria,
+            )
+        if isinstance(question, NoulQuestion):
+            yes = scores["true"]
+            no = scores["false"]
+            return NoulAnswer(type="noul", noul=1.0 / (1.0 + math.exp(no - yes)))
+        raise RuntimeErrorBase(f"unsupported question type: {type(question).__name__}")