lazycontroller 0.1.0__py3-none-any.whl

lazycontroller/__init__.py

@@ -0,0 +1,9 @@
+"""Drop-in ``openai`` clients: a fast decision model answers the structured calls it is sure of."""
+
+from .client import AsyncOpenAI, OpenAI
+from .lazy import FastConfig, Route
+from .shadow import ShadowReport
+
+__all__ = ["AsyncOpenAI", "FastConfig", "OpenAI", "Route", "ShadowReport"]
+
+__version__ = "0.1.0"

lazycontroller/build.py

@@ -0,0 +1,94 @@
+"""Build OpenAI response objects carrying the fast model's answer, indistinguishable to callers."""
+
+import json
+import time
+import uuid
+from collections.abc import Mapping
+from typing import Any
+
+# The SDK's own lenient constructor: a field added in a later SDK version must not break us.
+from openai._models import construct_type_unchecked
+from openai.types.chat import ChatCompletion
+from openai.types.responses import Response
+from typesafe_sdk import SystemOneResponse
+
+
+def chat_completion(output: Mapping[str, Any], answer: SystemOneResponse) -> ChatCompletion:
+    usage = _usage(answer)
+    return construct_type_unchecked(
+        type_=ChatCompletion,
+        value={
+            "id": f"chatcmpl-lazy-{uuid.uuid4().hex}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": answer.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "stop",
+                    "logprobs": None,
+                    "message": {
+                        "role": "assistant",
+                        "content": json.dumps(output),
+                        "refusal": None,
+                    },
+                }
+            ],
+            "usage": {
+                "prompt_tokens": usage[0],
+                "completion_tokens": usage[1],
+                "total_tokens": usage[0] + usage[1],
+            },
+        },
+    )
+
+
+def response(
+    output: Mapping[str, Any], answer: SystemOneResponse, kwargs: Mapping[str, Any]
+) -> Response:
+    usage = _usage(answer)
+    text = kwargs.get("text") or {}
+    return construct_type_unchecked(
+        type_=Response,
+        value={
+            "id": f"resp_lazy_{uuid.uuid4().hex}",
+            "object": "response",
+            "created_at": time.time(),
+            "status": "completed",
+            "model": answer.model,
+            "output": [
+                {
+                    "type": "message",
+                    "id": f"msg_lazy_{uuid.uuid4().hex}",
+                    "role": "assistant",
+                    "status": "completed",
+                    "content": [
+                        {
+                            "type": "output_text",
+                            "text": json.dumps(output),
+                            "annotations": [],
+                        }
+                    ],
+                }
+            ],
+            "instructions": kwargs.get("instructions"),
+            "metadata": kwargs.get("metadata"),
+            "text": text,
+            "tools": [],
+            "tool_choice": "auto",
+            "parallel_tool_calls": True,
+            "temperature": kwargs.get("temperature"),
+            "top_p": kwargs.get("top_p"),
+            "usage": {
+                "input_tokens": usage[0],
+                "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
+                "output_tokens": usage[1],
+                "output_tokens_details": {"reasoning_tokens": 0},
+                "total_tokens": usage[0] + usage[1],
+            },
+        },
+    )
+
+
+def _usage(answer: SystemOneResponse) -> tuple[int, int]:
+    return answer.usage.input_tokens or 0, answer.usage.output_tokens or 0

lazycontroller/client.py

@@ -0,0 +1,279 @@
+"""``OpenAI`` and ``AsyncOpenAI`` subclasses that let a fast model answer structured-output calls.
+
+Intercepted: ``chat.completions.create`` / ``parse`` and ``responses.create`` / ``parse``.
+Everything else is the stock SDK, and so is every intercepted call the fast model can't take.
+"""
+
+import os
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+import openai
+from openai.lib._parsing import parse_chat_completion, type_to_response_format_param
+from openai.lib._parsing._responses import parse_response, type_to_text_format_param
+from openai.resources.chat import AsyncChat, AsyncCompletions, Chat, Completions
+from openai.resources.responses import AsyncResponses, Responses
+
+from . import build
+from .lazy import (
+    AsyncRouter,
+    Call,
+    FastConfig,
+    Router,
+    chat_output_json,
+    responses_output_json,
+)
+from .requests import chat_request, given, responses_request
+from .schema import Ineligible
+
+
+def _chat_call(kwargs: dict[str, Any]) -> Call:
+    return Call(
+        "chat.completions",
+        kwargs.get("model"),
+        lambda: chat_request(kwargs),
+        lambda output, answer: build.chat_completion(output, answer),
+        chat_output_json,
+    )
+
+
+def _chat_parse_call(kwargs: dict[str, Any]) -> Call:
+    response_format = kwargs.get("response_format")
+
+    def prepare():
+        if not isinstance(response_format, type):
+            raise Ineligible("response_format is not a type")
+        return chat_request(
+            {
+                **kwargs,
+                "response_format": type_to_response_format_param(response_format),
+            }
+        )
+
+    def make(output, answer):
+        return parse_chat_completion(
+            response_format=response_format,
+            input_tools=kwargs.get("tools", openai.omit),
+            chat_completion=build.chat_completion(output, answer),
+        )
+
+    return Call("chat.completions", kwargs.get("model"), prepare, make, chat_output_json)
+
+
+def _responses_call(kwargs: dict[str, Any]) -> Call:
+    return Call(
+        "responses",
+        kwargs.get("model"),
+        lambda: responses_request(kwargs),
+        lambda output, answer: build.response(output, answer, given(kwargs)),
+        responses_output_json,
+    )
+
+
+def _responses_parse_call(kwargs: dict[str, Any]) -> Call:
+    text_format = kwargs.get("text_format")
+
+    def with_format() -> dict[str, Any]:
+        if not isinstance(text_format, type):
+            raise Ineligible("text_format is not a type")
+        if "parse_as" in given(kwargs):
+            raise Ineligible("uses parse_as")
+        rest = {k: v for k, v in kwargs.items() if k != "text_format"}
+        return {
+            **rest,
+            "text": {
+                **(rest.get("text") or {}),
+                "format": type_to_text_format_param(text_format),
+            },
+        }
+
+    def make(output, answer):
+        return parse_response(
+            text_format=text_format,
+            input_tools=kwargs.get("tools"),
+            response=build.response(output, answer, given(with_format())),
+        )
+
+    return Call(
+        "responses",
+        kwargs.get("model"),
+        lambda: responses_request(with_format()),
+        make,
+        responses_output_json,
+    )
+
+
+class _Completions(Completions):
+    # The overrides take **kwargs; type checkers keep seeing the SDK's full signatures.
+    if not TYPE_CHECKING:
+
+        def create(self, **kwargs):
+            return self._client._router.run(
+                _chat_call(kwargs), lambda: super(_Completions, self).create(**kwargs)
+            )
+
+        def parse(self, **kwargs):
+            return self._client._router.run(
+                _chat_parse_call(kwargs),
+                lambda: super(_Completions, self).parse(**kwargs),
+            )
+
+
+class _AsyncCompletions(AsyncCompletions):
+    if not TYPE_CHECKING:
+
+        async def create(self, **kwargs):
+            return await self._client._router.run(
+                _chat_call(kwargs),
+                lambda: super(_AsyncCompletions, self).create(**kwargs),
+            )
+
+        async def parse(self, **kwargs):
+            return await self._client._router.run(
+                _chat_parse_call(kwargs),
+                lambda: super(_AsyncCompletions, self).parse(**kwargs),
+            )
+
+
+class _Responses(Responses):
+    if not TYPE_CHECKING:
+
+        def create(self, **kwargs):
+            return self._client._router.run(
+                _responses_call(kwargs),
+                lambda: super(_Responses, self).create(**kwargs),
+            )
+
+        def parse(self, **kwargs):
+            return self._client._router.run(
+                _responses_parse_call(kwargs),
+                lambda: super(_Responses, self).parse(**kwargs),
+            )
+
+
+class _AsyncResponses(AsyncResponses):
+    if not TYPE_CHECKING:
+
+        async def create(self, **kwargs):
+            return await self._client._router.run(
+                _responses_call(kwargs),
+                lambda: super(_AsyncResponses, self).create(**kwargs),
+            )
+
+        async def parse(self, **kwargs):
+            return await self._client._router.run(
+                _responses_parse_call(kwargs),
+                lambda: super(_AsyncResponses, self).parse(**kwargs),
+            )
+
+
+class _Chat(Chat):
+    @cached_property
+    def completions(self) -> Completions:
+        return _Completions(self._client)
+
+
+class _AsyncChat(AsyncChat):
+    @cached_property
+    def completions(self) -> AsyncCompletions:
+        return _AsyncCompletions(self._client)
+
+
+_LOCAL_HOSTS = {"localhost", "127.0.0.1", "::1"}
+_CREDENTIALS = ("api_key", "admin_api_key", "workload_identity")
+
+
+def _with_local_key(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Fill in a placeholder key for a local server (Ollama, vLLM, LM Studio...).
+
+    The SDK refuses to start without credentials, but local servers ignore them. Only applies
+    where ``openai.OpenAI`` would have raised, so nothing changes for a configured client.
+    """
+    if any(kwargs.get(name) is not None for name in _CREDENTIALS):
+        return kwargs
+    if os.environ.get("OPENAI_API_KEY") or os.environ.get("OPENAI_ADMIN_KEY"):
+        return kwargs
+    base_url = kwargs.get("base_url") or os.environ.get("OPENAI_BASE_URL")
+    host = urlsplit(str(base_url)).hostname if base_url else None
+    if host in _LOCAL_HOSTS or (host or "").endswith(".localhost"):
+        return {**kwargs, "api_key": "local"}
+    return kwargs
+
+
+class OpenAI(openai.OpenAI):
+    """``openai.OpenAI`` that lets a fast model answer the structured-output calls it is sure of.
+
+    Takes every ``openai.OpenAI`` argument, plus ``fast`` for the fast model. A local
+    ``base_url`` such as Ollama's needs no ``api_key``.
+    """
+
+    _router: Router
+
+    def __init__(self, *args: Any, fast: FastConfig | None = None, **kwargs: Any) -> None:
+        super().__init__(*args, **_with_local_key(kwargs))
+        self._router = Router(fast or FastConfig())
+
+    @cached_property
+    def chat(self) -> Chat:
+        return _Chat(self)
+
+    @cached_property
+    def responses(self) -> Responses:
+        return _Responses(self)
+
+    def copy(self, *args: Any, **kwargs: Any) -> "OpenAI":
+        new = super().copy(*args, **kwargs)
+        new._router = self._router
+        return new
+
+    with_options = copy
+
+    def with_fast(self, **changes: Any) -> "OpenAI":
+        """A copy of this client with some ``FastConfig`` settings changed.
+
+        The ``with_options`` of routing: ``client.with_fast(min_confidence=0.95)`` for calls
+        where a wrong answer costs more, ``client.with_fast(enabled=False)`` to skip the fast
+        model.
+        The original client is unchanged.
+        """
+        new = self.copy()
+        new._router = self._router.derive(**changes)
+        return new
+
+
+class AsyncOpenAI(openai.AsyncOpenAI):
+    """``openai.AsyncOpenAI`` counterpart of ``lazycontroller.OpenAI``."""
+
+    _router: AsyncRouter
+
+    def __init__(self, *args: Any, fast: FastConfig | None = None, **kwargs: Any) -> None:
+        super().__init__(*args, **_with_local_key(kwargs))
+        self._router = AsyncRouter(fast or FastConfig())
+
+    @cached_property
+    def chat(self) -> AsyncChat:
+        return _AsyncChat(self)
+
+    @cached_property
+    def responses(self) -> AsyncResponses:
+        return _AsyncResponses(self)
+
+    def copy(self, *args: Any, **kwargs: Any) -> "AsyncOpenAI":
+        new = super().copy(*args, **kwargs)
+        new._router = self._router
+        return new
+
+    with_options = copy
+
+    def with_fast(self, **changes: Any) -> "AsyncOpenAI":
+        """A copy of this client with some ``FastConfig`` settings changed.
+
+        The ``with_options`` of routing: ``client.with_fast(min_confidence=0.95)`` for calls
+        where a wrong answer costs more, ``client.with_fast(enabled=False)`` to skip the fast
+        model.
+        The original client is unchanged.
+        """
+        new = self.copy()
+        new._router = self._router.derive(**changes)
+        return new

lazycontroller/lazy.py

@@ -0,0 +1,265 @@
+"""Decide, per request, whether the fast model answers or the call goes on to the slow one.
+
+The fast model is a decision model on a System One API: Jev on TypeSafe, a local Kev on
+Ollama, or the same models hosted elsewhere. The slow model is whatever the caller's
+OpenAI-compatible client was going to call anyway.
+"""
+
+import json
+import logging
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass, field, replace
+from typing import Any, Literal, Self
+
+from typesafe_sdk import (
+    AsyncTypeSafeClient,
+    SystemOneResponse,
+    TypeSafeClient,
+)
+
+from .requests import Request
+from .schema import Ineligible, Plan, compile_schema
+
+logger = logging.getLogger("lazycontroller")
+
+type Api = Literal["chat.completions", "responses"]
+type Reason = Literal[
+    "confident", "low_confidence", "ineligible", "fast_model_error", "disabled", "shadow"
+]
+
+
+@dataclass(frozen=True)
+class Route:
+    """How one request was answered; passed to ``FastConfig.on_route``."""
+
+    api: Api
+    answered_by: Literal["fast", "slow"]
+    reason: Reason
+    """Why: the fast model was ``confident``, or the slow model was called because of
+    ``low_confidence``, an ``ineligible`` request, a ``fast_model_error``, routing being
+    ``disabled``, or ``shadow`` mode."""
+    slow_model: str | None
+    """The model the caller asked for."""
+    fast_model: str | None = None
+    confidence: float | None = None
+    """The lowest field confidence, which is what ``min_confidence`` is compared with."""
+    field_confidences: Mapping[str, float] = field(default_factory=dict)
+    fast_output: Mapping[str, Any] | None = None
+    """The fast model's answer as a JSON object, whenever it was asked (also when the slow model
+    answered in the end, so shadow disagreements can be inspected)."""
+    fast_input_tokens: int | None = None
+    detail: str | None = None
+    """Why the request was ineligible, or the fast model's error."""
+    shadow_mismatches: tuple[str, ...] | None = None
+    """In shadow mode, the fields where the two models disagreed (empty when they agree)."""
+
+
+@dataclass(frozen=True)
+class FastConfig:
+    """The fast decision model, and when its answer is used instead of the slow model's."""
+
+    base_url: str | None = None
+    """The System One server: ``https://api.typesafe.ai``, ``https://openrouter.ai/api``,
+    ``http://localhost:11434`` for Ollama, ... Without it (or ``client``) the fast side is off
+    and every call goes to the slow model."""
+    api_key: str | None = None
+    """Key for ``base_url``, and sent nowhere else. Local servers need none."""
+    model: str | None = None
+    """The decision model; ``None`` means the server's default."""
+    min_confidence: float = 0.9
+    """Every field's confidence must reach this for the fast model's answer to be used;
+    otherwise the call goes on to the slow model."""
+    enabled: bool = True
+    """``False`` sends every call straight to the slow model."""
+    shadow: bool = False
+    """Ask the fast model but always return the slow model's answer, reporting where they
+    differ."""
+    on_route: Callable[[Route], None] | None = None
+    client: TypeSafeClient | AsyncTypeSafeClient | None = None
+    """Use this System One client instead of creating one."""
+
+    def __post_init__(self) -> None:
+        if not 0.0 <= self.min_confidence <= 1.0:
+            raise ValueError("min_confidence must be between 0 and 1")
+
+
+class _Router[C]:
+    def __init__(self, config: FastConfig, client_cls: type[C]) -> None:
+        self.config = config
+        self._client_cls = client_cls
+        self._fast: C | None = config.client  # type: ignore[assignment]
+
+    @property
+    def fast(self) -> C | None:
+        if self._fast is None:
+            if base_url := self.config.base_url:
+                # Passed explicitly, so the SDK reads neither TYPESAFE_API_KEY nor
+                # TYPESAFE_BASE_URL. It insists on some key; local servers ignore it.
+                api_key = self.config.api_key or "local"
+                self._fast = self._client_cls(api_key=api_key, base_url=base_url)
+        return self._fast
+
+    def derive(self, **changes: Any) -> Self:
+        """A router with some settings changed, sharing the fast model's connection when it can."""
+        config = replace(self.config, **changes)
+        router = type(self)(config)
+        if all(
+            getattr(config, name) == getattr(self.config, name)
+            for name in ("api_key", "base_url", "client")
+        ):
+            router._fast = self._fast
+        return router
+
+    def plan(
+        self, api: Api, prepare: Callable[[], Request], model: str | None
+    ) -> tuple[Request, Plan] | Route:
+        if not self.config.enabled:
+            return Route(api, "slow", "disabled", model, detail="disabled in FastConfig")
+        try:
+            if self.fast is None:
+                return Route(api, "slow", "disabled", model, detail="no base_url in FastConfig")
+            request = prepare()
+            return request, compile_schema(request.schema, title=request.name)
+        except Ineligible as error:
+            return Route(api, "slow", "ineligible", model, detail=str(error))
+        except Exception as error:
+            # A request we can't read is the slow model's to accept or reject, not ours.
+            logger.warning("Could not read request, passing it to the slow model: %r", error)
+            detail = f"could not read request: {type(error).__name__}: {error}"
+            return Route(api, "slow", "ineligible", model, detail=detail)
+
+    def judge(
+        self, api: Api, plan: Plan, answer: SystemOneResponse, model: str | None
+    ) -> tuple[dict[str, Any], Route]:
+        output, confidences = plan.assemble(answer.answers)
+        confidence = min(confidences.values())
+        confident = confidence >= self.config.min_confidence
+        route = Route(
+            api,
+            "fast" if confident and not self.config.shadow else "slow",
+            "shadow" if self.config.shadow else "confident" if confident else "low_confidence",
+            model,
+            fast_model=answer.model,
+            confidence=confidence,
+            field_confidences=confidences,
+            fast_output=output,
+            fast_input_tokens=answer.usage.input_tokens,
+        )
+        return output, route
+
+    def fast_failed(self, api: Api, model: str | None, error: Exception) -> Route:
+        logger.warning("Fast model failed, falling back to the slow model: %r", error)
+        detail = f"{type(error).__name__}: {error}"
+        return Route(api, "slow", "fast_model_error", model, detail=detail)
+
+    def report(self, route: Route) -> None:
+        if self.config.on_route is not None:
+            self.config.on_route(route)
+
+
+@dataclass(frozen=True)
+class Call:
+    """One intercepted SDK call: how to read it, how to answer it, how to pass it on."""
+
+    api: Api
+    model: str | None
+    prepare: Callable[[], Request]
+    build: Callable[[dict[str, Any], SystemOneResponse], Any]
+    output_json: Callable[[Any], Any]
+    """Extract the JSON object from a slow-model response, for shadow comparisons."""
+
+
+class Router(_Router[TypeSafeClient]):
+    def __init__(self, config: FastConfig) -> None:
+        super().__init__(config, TypeSafeClient)
+
+    def run(self, call: Call, slow: Callable[[], Any]) -> Any:
+        planned = self.plan(call.api, call.prepare, call.model)
+        if isinstance(planned, Route):
+            self.report(planned)
+            return slow()
+        request, plan = planned
+        try:
+            answer = self.fast.system_one(  # type: ignore[union-attr]
+                state=request.state, questions=plan.questions, model=self.config.model
+            )
+            output, route = self.judge(call.api, plan, answer, call.model)
+        except Exception as error:
+            self.report(self.fast_failed(call.api, call.model, error))
+            return slow()
+        if route.answered_by == "fast":
+            try:
+                result = call.build(output, answer)
+            except Exception as error:  # e.g. a validator in the caller's model rejects it
+                self.report(self.fast_failed(call.api, call.model, error))
+                return slow()
+            self.report(route)
+            return result
+        result = slow()
+        self.report(_with_shadow(route, plan, output, call, result))
+        return result
+
+
+class AsyncRouter(_Router[AsyncTypeSafeClient]):
+    def __init__(self, config: FastConfig) -> None:
+        super().__init__(config, AsyncTypeSafeClient)
+
+    async def run(self, call: Call, slow: Callable[[], Awaitable[Any]]) -> Any:
+        planned = self.plan(call.api, call.prepare, call.model)
+        if isinstance(planned, Route):
+            self.report(planned)
+            return await slow()
+        request, plan = planned
+        try:
+            answer = await self.fast.system_one(  # type: ignore[union-attr]
+                state=request.state, questions=plan.questions, model=self.config.model
+            )
+            output, route = self.judge(call.api, plan, answer, call.model)
+        except Exception as error:
+            self.report(self.fast_failed(call.api, call.model, error))
+            return await slow()
+        if route.answered_by == "fast":
+            try:
+                result = call.build(output, answer)
+            except Exception as error:  # e.g. a validator in the caller's model rejects it
+                self.report(self.fast_failed(call.api, call.model, error))
+                return await slow()
+            self.report(route)
+            return result
+        result = await slow()
+        self.report(_with_shadow(route, plan, output, call, result))
+        return result
+
+
+def _with_shadow(
+    route: Route, plan: Plan, output: dict[str, Any], call: Call, result: Any
+) -> Route:
+    if route.reason != "shadow":
+        return route
+    try:
+        theirs = call.output_json(result)
+    except (TypeError, ValueError, AttributeError, IndexError):
+        theirs = None
+    mismatches = tuple(
+        ".".join(f.path) for f in plan.fields if _get(output, f.path) != _get(theirs, f.path)
+    )
+    return Route(**{**route.__dict__, "shadow_mismatches": mismatches})
+
+
+def _get(obj: Any, path: tuple[str, ...]) -> Any:
+    for key in path:
+        if not isinstance(obj, Mapping) or key not in obj:
+            return _MISSING
+        obj = obj[key]
+    return obj
+
+
+_MISSING = object()
+
+
+def chat_output_json(completion: Any) -> Any:
+    return json.loads(completion.choices[0].message.content)
+
+
+def responses_output_json(response: Any) -> Any:
+    return json.loads(response.output_text)