perf: reuse serializer wrapper models instead of rebuilding them on every call (#1335)

Author: jason-schrader-openrouter

src/openrouter/_hooks/registration.py

@@ -1,4 +1,5 @@
 from .types import Hooks
+from . import serializer_cache
 
 
 # This file is only ever generated once on the first generation and then is free to be modified.
@@ -11,3 +12,4 @@ def init_hooks(hooks: Hooks):
     """Add hooks by calling hooks.register{sdk_init/before_request/after_success/after_error}Hook
     with an instance of a hook that implements that specific Hook interface
     Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance"""
+    serializer_cache.install()

src/openrouter/_hooks/serializer_cache.py

@@ -0,0 +1,63 @@
+"""Reuse the wrapper models that ``utils.serializers`` builds for each type.
+
+``marshal_json`` and ``unmarshal`` call ``pydantic.create_model`` on every
+call, so each request builds and compiles new model classes. This module
+replaces ``serializers.create_model`` with a version that builds one model per
+(name, type, config) and returns it on later calls. Calls in any other shape,
+or with unhashable arguments, fall through to ``pydantic.create_model``.
+"""
+
+import functools
+from typing import Any, Callable, Tuple, Type, get_args, get_origin
+
+from pydantic import BaseModel, ConfigDict
+
+from ..utils import serializers
+
+_uncached_create_model: Callable[..., Type[BaseModel]] = serializers.create_model
+
+# Bounded so dynamically created types can't grow the cache without limit.
+_CACHE_SIZE = 1024
+
+
+def _type_key(typ: Any) -> Any:
+    # Equal annotations can still validate differently: `int | float ==
+    # float | int` and `Literal[1] == Literal[True]`, but pydantic tries union
+    # members in order and checks literal types. Keep order and leaf types.
+    args = get_args(typ)
+    if not args:
+        return (type(typ), typ)
+    return (get_origin(typ), tuple(_type_key(arg) for arg in args))
+
+
+@functools.lru_cache(maxsize=_CACHE_SIZE)
+def _cached_wrapper_model(
+    name: str, type_key: Any, config_items: Tuple[Tuple[str, Any], ...], typ: Any
+) -> Type[BaseModel]:
+    del type_key
+    return _uncached_create_model(
+        name, body=(typ, ...), __config__=ConfigDict(**dict(config_items))  # type: ignore[typeddict-item]
+    )
+
+
+def _create_model(name: str, **kwargs: Any) -> Type[BaseModel]:
+    body, config = kwargs.get("body"), kwargs.get("__config__")
+    if (
+        set(kwargs) != {"body", "__config__"}
+        or not isinstance(body, tuple)
+        or len(body) != 2
+        or body[1] is not ...
+        or not isinstance(config, dict)
+    ):
+        return _uncached_create_model(name, **kwargs)
+    typ = body[0]
+    try:
+        return _cached_wrapper_model(
+            name, _type_key(typ), tuple(sorted(config.items())), typ
+        )
+    except TypeError:
+        return _uncached_create_model(name, **kwargs)
+
+
+def install() -> None:
+    setattr(serializers, "create_model", _create_model)

tests/test_serializer_cache.py

@@ -0,0 +1,160 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+import httpx
+import pydantic
+import pytest
+
+from openrouter import OpenRouter
+from openrouter._hooks import serializer_cache
+from openrouter.utils import serializers
+
+COMPLETION = {
+    "id": "gen-test",
+    "object": "chat.completion",
+    "created": 1760000000,
+    "model": "openai/gpt-4o-mini",
+    "system_fingerprint": None,
+    "choices": [
+        {
+            "index": 0,
+            "finish_reason": "stop",
+            "message": {"role": "assistant", "content": "ok"},
+        }
+    ],
+    "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+}
+
+
+def _client() -> OpenRouter:
+    transport = httpx.MockTransport(
+        lambda request: httpx.Response(200, json=COMPLETION)
+    )
+    return OpenRouter(api_key="test-key", client=httpx.Client(transport=transport))
+
+
+def _send(client: OpenRouter):
+    return client.chat.send(
+        model="openai/gpt-4o-mini",
+        messages=[{"role": "user", "content": "hi"}],
+    )
+
+
+def test_sdk_init_installs_serializer_cache():
+    _client()
+
+    assert serializers.create_model is serializer_cache._create_model
+
+
+def test_repeated_requests_reuse_serializer_models(monkeypatch):
+    built = []
+    uncached_create_model = serializer_cache._uncached_create_model
+
+    def counting_create_model(*args, **kwargs):
+        built.append(args[0] if args else kwargs.get("__model_name"))
+        return uncached_create_model(*args, **kwargs)
+
+    client = _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+    monkeypatch.setattr(serializer_cache, "_uncached_create_model", counting_create_model)
+
+    first = _send(client)
+    built_by_first_request = len(built)
+    second = _send(client)
+
+    assert built_by_first_request > 0
+    assert len(built) == built_by_first_request
+    assert first.choices[0].message.content == "ok"
+    assert second.model_dump() == first.model_dump()
+
+
+def _uncached(monkeypatch):
+    monkeypatch.setattr(
+        serializers, "create_model", serializer_cache._uncached_create_model
+    )
+
+
+def _outcome(fn):
+    try:
+        result = fn()
+    except pydantic.ValidationError as error:
+        return ("error", [(e["type"], e["loc"]) for e in error.errors()])
+    return ("ok", type(result), result)
+
+
+class _Aliased(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(populate_by_name=True)
+    max_tokens: Optional[int] = pydantic.Field(default=None, alias="maxTokens")
+
+
+CASES = [
+    ("1", int | float),
+    ("1", float | int),
+    ("1", Union[int, float]),
+    ("1", Union[float, int]),
+    (True, Literal[1]),
+    (True, Literal[True]),
+    (1, Literal[True]),
+    ({"a": [1, 2]}, dict[str, list[int]]),
+    ({"a": ["x"]}, dict[str, list[int]]),
+    ({"maxTokens": 3}, _Aliased),
+    ({"max_tokens": 3}, _Aliased),
+    ({"maxTokens": "x"}, _Aliased),
+    (None, Optional[_Aliased]),
+    ([1, "2"], List[int]),
+]
+
+
+def _all_outcomes():
+    unmarshalled = [_outcome(lambda: serializers.unmarshal(v, t)) for v, t in CASES]
+    marshalled = [
+        _outcome(lambda: serializers.marshal_json(v, t)) for v, t in CASES
+    ]
+    return unmarshalled + marshalled
+
+
+@pytest.mark.parametrize("order", [1, -1], ids=["forward", "reverse"])
+def test_cached_models_behave_like_uncached_ones(monkeypatch, order):
+    with monkeypatch.context() as m:
+        _uncached(m)
+        expected = _all_outcomes()
+
+    _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+    # Warm the cache in the given order, then compare every case.
+    for value, typ in CASES[::order]:
+        _outcome(lambda: serializers.unmarshal(value, typ))
+        _outcome(lambda: serializers.marshal_json(value, typ))
+
+    assert _all_outcomes() == expected
+
+
+def test_unhashable_types_fall_back_to_uncached_models(monkeypatch):
+    built = []
+    uncached_create_model = serializer_cache._uncached_create_model
+
+    def counting_create_model(*args, **kwargs):
+        built.append(args[0])
+        return uncached_create_model(*args, **kwargs)
+
+    _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+    monkeypatch.setattr(serializer_cache, "_uncached_create_model", counting_create_model)
+    typ = Annotated[int, {"unhashable": []}]
+
+    assert serializers.unmarshal("5", typ) == 5
+    assert serializers.unmarshal("5", typ) == 5
+    assert built == ["Unmarshaller", "Unmarshaller"]
+    assert serializer_cache._cached_wrapper_model.cache_info().currsize == 0
+
+
+def test_cache_size_is_bounded():
+    _client()
+    serializer_cache._cached_wrapper_model.cache_clear()
+
+    for i in range(serializer_cache._CACHE_SIZE + 50):
+        dynamic = pydantic.create_model(f"Dynamic{i}", x=(int, ...))
+        serializers.unmarshal({"x": i}, dynamic)
+
+    info = serializer_cache._cached_wrapper_model.cache_info()
+    assert info.maxsize == serializer_cache._CACHE_SIZE
+    assert info.currsize == serializer_cache._CACHE_SIZE