update

libs/sdk-py/langgraph_sdk/cache.py

@@ -7,9 +7,18 @@ Values must be JSON-serializable (dicts, lists, strings, numbers, booleans,
 
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
-from datetime import timedelta
-from typing import Any, Generic, Literal, TypeVar
+import dataclasses
+import enum
+import functools
+import inspect
+from collections.abc import Awaitable, Callable, Mapping
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from pathlib import PurePath
+from typing import Any, Generic, Literal, TypeVar, get_type_hints, overload
+from uuid import UUID
+
+import orjson
 
 T = TypeVar("T")
 
@@ -50,6 +59,7 @@ __all__ = [
     "cache_get",
     "cache_set",
     "swr",
+    "swr_cached",
 ]
 
 
@@ -138,3 +148,268 @@ async def swr(
     return await _api_swr(
         key, loader, fresh_for=fresh_for, max_age=max_age, model=model
     )
+
+
+def _build_cache_key(
+    module: str,
+    qualname: str,
+    sig: inspect.Signature,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> str:
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    payload = {
+        "module": module,
+        "qualname": qualname,
+        "args": [
+            {
+                "name": name,
+                "value": _normalize_cache_key_value(value, name=name, path=name),
+            }
+            for name, value in bound.arguments.items()
+        ],
+    }
+    return orjson.dumps(payload, option=orjson.OPT_SORT_KEYS).decode()
+
+
+def _type_identifier(tp: type[Any]) -> str:
+    return f"{tp.__module__}.{tp.__qualname__}"
+
+
+def _stable_key_dump(value: Any) -> bytes:
+    return orjson.dumps(value, option=orjson.OPT_SORT_KEYS)
+
+
+def _normalize_cache_key_value(
+    value: Any,
+    *,
+    name: str | None = None,
+    path: str = "value",
+) -> Any:
+    if name in {"self", "cls"}:
+        cls = value if isinstance(value, type) else type(value)
+        return {"class": _type_identifier(cls), "kind": name}
+
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+
+    if isinstance(value, bytes):
+        return {"hex": value.hex(), "kind": "bytes"}
+
+    if isinstance(value, (datetime, date, time)):
+        return {"kind": type(value).__name__, "value": value.isoformat()}
+
+    if isinstance(value, timedelta):
+        return {"kind": "timedelta", "value": value.total_seconds()}
+
+    if isinstance(value, Decimal):
+        return {"kind": "decimal", "value": str(value)}
+
+    if isinstance(value, UUID):
+        return {"kind": "uuid", "value": str(value)}
+
+    if isinstance(value, PurePath):
+        return {"kind": "path", "value": str(value)}
+
+    if isinstance(value, enum.Enum):
+        return {
+            "kind": "enum",
+            "type": _type_identifier(type(value)),
+            "value": _normalize_cache_key_value(value.value, path=f"{path}.value"),
+        }
+
+    if dataclasses.is_dataclass(value) and not isinstance(value, type):
+        return {
+            "fields": {
+                field.name: _normalize_cache_key_value(
+                    getattr(value, field.name),
+                    path=f"{path}.{field.name}",
+                )
+                for field in dataclasses.fields(value)
+            },
+            "kind": "dataclass",
+            "type": _type_identifier(type(value)),
+        }
+
+    if hasattr(value, "model_dump") and callable(value.model_dump):
+        if isinstance(value, type):
+            raise TypeError(
+                f"Cannot auto-generate a stable cache key for `{path}` from the "
+                f"type object {value!r}. Pass `key=` to `swr_cached` instead."
+            )
+        return {
+            "kind": "model",
+            "type": _type_identifier(type(value)),
+            "value": _normalize_cache_key_value(
+                value.model_dump(mode="json"),
+                path=path,
+            ),
+        }
+
+    if hasattr(value, "dict") and callable(value.dict):
+        if isinstance(value, type):
+            raise TypeError(
+                f"Cannot auto-generate a stable cache key for `{path}` from the "
+                f"type object {value!r}. Pass `key=` to `swr_cached` instead."
+            )
+        return {
+            "kind": "model",
+            "type": _type_identifier(type(value)),
+            "value": _normalize_cache_key_value(value.dict(), path=path),
+        }
+
+    if isinstance(value, Mapping):
+        items = [
+            {
+                "key": _normalize_cache_key_value(key, path=f"{path}.<key>"),
+                "value": _normalize_cache_key_value(
+                    item,
+                    path=f"{path}[{key!r}]",
+                ),
+            }
+            for key, item in value.items()
+        ]
+        items.sort(key=lambda item: _stable_key_dump(item["key"]))
+        return {"items": items, "kind": "mapping"}
+
+    if isinstance(value, tuple):
+        return {
+            "items": [
+                _normalize_cache_key_value(item, path=f"{path}[{index}]")
+                for index, item in enumerate(value)
+            ],
+            "kind": "tuple",
+        }
+
+    if isinstance(value, list):
+        return {
+            "items": [
+                _normalize_cache_key_value(item, path=f"{path}[{index}]")
+                for index, item in enumerate(value)
+            ],
+            "kind": "list",
+        }
+
+    if isinstance(value, (set, frozenset)):
+        items = [_normalize_cache_key_value(item, path=f"{path}[]") for item in value]
+        items.sort(key=_stable_key_dump)
+        return {"items": items, "kind": type(value).__name__}
+
+    if isinstance(value, type):
+        return {"kind": "type", "type": _type_identifier(value)}
+
+    if type(value).__repr__ is object.__repr__:
+        raise TypeError(
+            f"Cannot auto-generate a stable cache key for `{path}` of type "
+            f"`{_type_identifier(type(value))}`. Pass `key=` to `swr_cached` "
+            "instead."
+        )
+
+    return {
+        "kind": "repr",
+        "repr": repr(value),
+        "type": _type_identifier(type(value)),
+    }
+
+
+def _get_model_from_hints(func: Callable[..., Any]) -> type | None:
+    try:
+        hints = get_type_hints(func)
+    except Exception:
+        return None
+    ret = hints.get("return")
+    if ret is None:
+        return None
+    try:
+        from pydantic import BaseModel
+    except ImportError:
+        return None
+    if isinstance(ret, type) and issubclass(ret, BaseModel):
+        return ret
+    return None
+
+
+@overload
+def swr_cached(
+    fn: Callable[..., Awaitable[T]],
+    /,
+) -> Callable[..., Awaitable[SWRResult[T]]]: ...
+
+
+@overload
+def swr_cached(
+    *,
+    key: str | Callable[..., str] | None = ...,
+    fresh_for: timedelta | None = ...,
+    max_age: timedelta | None = ...,
+    model: type[T] | None = ...,
+) -> Callable[
+    [Callable[..., Awaitable[T]]], Callable[..., Awaitable[SWRResult[T]]]
+]: ...
+
+
+def swr_cached(
+    fn=None,
+    *,
+    key=None,
+    fresh_for=None,
+    max_age=None,
+    model=None,
+):
+    """Decorator that wraps an async function with :func:`swr` caching.
+
+    Can be used with or without parentheses::
+
+        @swr_cached
+        async def fetch_config():
+            ...
+
+        @swr_cached(fresh_for=timedelta(minutes=5))
+        async def fetch_profile(user_id: str) -> Profile:
+            ...
+
+    The cache key is auto-derived from the function's module, qualified name,
+    and a structured serialization of the bound call arguments. For methods,
+    ``self`` and ``cls`` are keyed by class identity rather than object
+    instance identity. Override with ``key=`` (a static string or a callable
+    that receives the same arguments as the decorated function) when method
+    state matters or arguments are not stably serializable.
+
+    If the return annotation is a Pydantic `BaseModel` subclass and
+    ``model`` is not provided, the model is inferred automatically.
+    """
+
+    def decorator(
+        func: Callable[..., Awaitable[T]],
+    ) -> Callable[..., Awaitable[SWRResult[T]]]:
+        sig = inspect.signature(func)
+        resolved_model = model
+        if resolved_model is None:
+            resolved_model = _get_model_from_hints(func)
+
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> SWRResult[T]:
+            if key is None:
+                module = getattr(func, "__module__", None) or "unknown"
+                qualname = getattr(func, "__qualname__", None) or getattr(
+                    func, "__name__", "unknown"
+                )
+                cache_key = _build_cache_key(module, qualname, sig, args, kwargs)
+            elif callable(key):
+                cache_key = key(*args, **kwargs)
+            else:
+                cache_key = key
+            return await swr(
+                cache_key,
+                lambda: func(*args, **kwargs),
+                fresh_for=fresh_for,
+                max_age=max_age,
+                model=resolved_model,
+            )
+
+        return wrapper
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator

libs/sdk-py/tests/test_cache.py

@@ -1,9 +1,12 @@
+import inspect
 from datetime import timedelta
 from unittest.mock import AsyncMock
 
+import orjson
 import pytest
 
 import langgraph_sdk.cache as cache_module
+from langgraph_sdk.cache import swr_cached
 
 
 @pytest.mark.asyncio
@@ -67,3 +70,239 @@ async def test_swr_defaults(monkeypatch):
 
     assert forwarded["fresh_for"] == timedelta(0)
     assert forwarded["max_age"] == timedelta(days=1)
+
+
+# -- swr_cached decorator tests --
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_no_parens(monkeypatch):
+    """`@swr_cached` without parentheses uses a structured auto-derived key."""
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["key"] = key
+        forwarded["model"] = model
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached
+    async def fetch_config():
+        return {"debug": True}
+
+    result = await fetch_config()
+    cache_key = orjson.loads(forwarded["key"])
+
+    assert result == {"debug": True}
+    assert cache_key == {
+        "args": [],
+        "module": fetch_config.__module__,
+        "qualname": fetch_config.__qualname__,
+    }
+    assert forwarded["model"] is None
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_with_options(monkeypatch):
+    """@swr_cached(...) with keyword options."""
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["key"] = key
+        forwarded["fresh_for"] = fresh_for
+        forwarded["max_age"] = max_age
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached(fresh_for=timedelta(minutes=5), max_age=timedelta(hours=1))
+    async def fetch_config():
+        return "ok"
+
+    await fetch_config()
+
+    assert forwarded["fresh_for"] == timedelta(minutes=5)
+    assert forwarded["max_age"] == timedelta(hours=1)
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_key_includes_args(monkeypatch):
+    """Arguments are serialized structurally into the auto-derived cache key."""
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["key"] = key
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached
+    async def fetch_profile(user_id: str):
+        return {"id": user_id}
+
+    await fetch_profile("abc123")
+    cache_key = orjson.loads(forwarded["key"])
+    assert cache_key["args"] == [{"name": "user_id", "value": "abc123"}]
+
+
+def test_build_cache_key_distinguishes_modules():
+    async def fetch_profile(user_id: str): ...
+
+    sig = inspect.signature(fetch_profile)
+    key_one = cache_module._build_cache_key(
+        "alpha.module", "fetch_profile", sig, ("1",), {}
+    )
+    key_two = cache_module._build_cache_key(
+        "beta.module", "fetch_profile", sig, ("1",), {}
+    )
+
+    assert key_one != key_two
+
+
+def test_build_cache_key_distinguishes_argument_boundaries():
+    async def fetch_profile(left: str, right: str): ...
+
+    sig = inspect.signature(fetch_profile)
+    key_one = cache_module._build_cache_key(
+        "alpha.module",
+        "fetch_profile",
+        sig,
+        ("a:b", "c"),
+        {},
+    )
+    key_two = cache_module._build_cache_key(
+        "alpha.module",
+        "fetch_profile",
+        sig,
+        ("a", "b:c"),
+        {},
+    )
+
+    assert key_one != key_two
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_explicit_key_string(monkeypatch):
+    """Explicit string key overrides auto-derivation."""
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["key"] = key
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached(key="my-static-key")
+    async def fetch_stuff():
+        return 42
+
+    await fetch_stuff()
+    assert forwarded["key"] == "my-static-key"
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_explicit_key_callable(monkeypatch):
+    """Explicit callable key receives the function's arguments."""
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["key"] = key
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached(key=lambda org, repo: f"repo:{org}/{repo}")
+    async def fetch_repo(org: str, repo: str):
+        return {"full_name": f"{org}/{repo}"}
+
+    await fetch_repo("langchain-ai", "langgraph")
+    assert forwarded["key"] == "repo:langchain-ai/langgraph"
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_method_key_uses_class_identity(monkeypatch):
+    forwarded = []
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded.append(key)
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    class Client:
+        @swr_cached
+        async def fetch_repo(self, repo: str):
+            return {"repo": repo}
+
+    await Client().fetch_repo("langgraph")
+    await Client().fetch_repo("langgraph")
+
+    assert forwarded[0] == forwarded[1]
+    cache_key = orjson.loads(forwarded[0])
+    assert cache_key["args"][0] == {
+        "name": "self",
+        "value": {
+            "class": f"{Client.__module__}.{Client.__qualname__}",
+            "kind": "self",
+        },
+    }
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_rejects_unstable_object_args():
+    class Unstable:
+        pass
+
+    @swr_cached
+    async def fetch_profile(user: Unstable):
+        return {"user_type": type(user).__name__}
+
+    with pytest.raises(
+        TypeError,
+        match="Cannot auto-generate a stable cache key for `user`",
+    ):
+        await fetch_profile(Unstable())
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_preserves_function_metadata(monkeypatch):
+    """functools.wraps preserves __name__ and __doc__."""
+
+    async def fake_swr(_key, loader, **_kw):
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    @swr_cached
+    async def my_loader():
+        """My docstring."""
+        return 1
+
+    assert my_loader.__name__ == "my_loader"
+    assert my_loader.__doc__ == "My docstring."
+
+
+@pytest.mark.asyncio
+async def test_swr_cached_infers_pydantic_model(monkeypatch):
+    """Model is auto-detected from return type annotation."""
+    pytest.importorskip("pydantic")
+    from pydantic import BaseModel
+
+    forwarded = {}
+
+    async def fake_swr(key, loader, *, fresh_for, max_age, model):  # noqa: ARG001
+        forwarded["model"] = model
+        return await loader()
+
+    monkeypatch.setattr(cache_module, "_api_swr", fake_swr)
+
+    class Profile(BaseModel):
+        name: str
+
+    @swr_cached
+    async def fetch_profile() -> Profile:
+        return Profile(name="Alice")
+
+    await fetch_profile()
+    assert forwarded["model"] is Profile