Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/openrouter/_hooks/registration.py
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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()
46 changes: 46 additions & 0 deletions src/openrouter/_hooks/serializer_cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
"""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

from pydantic import BaseModel, ConfigDict

from ..utils import serializers

_uncached_create_model: Callable[..., Type[BaseModel]] = serializers.create_model


@functools.lru_cache(maxsize=None)
def _cached_wrapper_model(
name: str, typ: Any, config_items: Tuple[Tuple[str, Any], ...]
) -> Type[BaseModel]:
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)
try:
return _cached_wrapper_model(name, body[0], tuple(sorted(config.items())))
except TypeError:
return _uncached_create_model(name, **kwargs)


def install() -> None:
setattr(serializers, "create_model", _create_model)
76 changes: 76 additions & 0 deletions tests/test_serializer_cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
import json

import httpx

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 test_cached_models_serialize_like_uncached_ones():
_client()

value = {"a": [1, 2]}
typ = dict[str, list[int]]

assert json.loads(serializers.marshal_json(value, typ)) == value
assert serializers.unmarshal(value, typ) == value
assert serializers.unmarshal(value, typ) == value
Loading