Separate immutable prediction cache keys from mutable call arguments

great_ai/deploy/great_ai.py

@@ -23,6 +23,7 @@ from ..constants import DASHBOARD_PATH
 from ..context import get_context
 from ..external.async_lru import alru_cache
 from ..helper import freeze_arguments, get_function_metadata_store, snake_case_to_text
+from ..helper.frozen_call import FrozenCall
 from ..models.use_model import model_versions
 from ..parameters.automatically_decorate_parameters import (
     automatically_decorate_parameters,
@@ -258,38 +259,54 @@ class GreatAI(Generic[T, V]):
         func: Callable[..., Union[V, Awaitable[V]]],
     ) -> Callable[..., T]:
         @lru_cache(maxsize=get_context().prediction_cache_size)
-        def func_in_tracing_context_sync(
-            *args: Any,
-            do_not_persist_traces: bool = False,
-            **kwargs: Any,
-        ) -> T:
+        def func_in_tracing_context_sync(call: FrozenCall) -> T:
             with TracingContext[V](
-                func.__name__, do_not_persist_traces=do_not_persist_traces
+                func.__name__, do_not_persist_traces=call.do_not_persist
             ) as t:
-                result = func(*args, **kwargs)
+                result = func(*call.args, **call.kwargs)
                 return cast(T, t.finalise(output=result))
 
         @alru_cache(maxsize=get_context().prediction_cache_size, cache_exceptions=False)
-        async def func_in_tracing_context_async(
-            *args: Any,
-            do_not_persist_traces: bool = False,
-            **kwargs: Any,
-        ) -> T:
+        async def func_in_tracing_context_async(call: FrozenCall) -> T:
             with TracingContext[V](
-                func.__name__, do_not_persist_traces=do_not_persist_traces
+                func.__name__, do_not_persist_traces=call.do_not_persist
             ) as t:
-                result = await cast(Callable[..., Awaitable], func)(*args, **kwargs)
+                result = await cast(Callable[..., Awaitable], func)(
+                    *call.args, **call.kwargs
+                )
                 return cast(T, t.finalise(output=result))
 
-        return cast(
-            Callable[..., T],
-            (
-                func_in_tracing_context_async
-                if get_function_metadata_store(func).is_asynchronous
-                else func_in_tracing_context_sync
-            ),
+        @wraps(func)
+        def sync(*args: Any, do_not_persist_traces: bool = False, **kwargs: Any) -> T:
+            return func_in_tracing_context_sync(
+                FrozenCall(args, kwargs, do_not_persist_traces)
             )
 
+        @wraps(func)
+        async def asynchronous(
+            *args: Any, do_not_persist_traces: bool = False, **kwargs: Any
+        ) -> T:
+            return await func_in_tracing_context_async(
+                FrozenCall(args, kwargs, do_not_persist_traces)
+            )
+
+        is_async = get_function_metadata_store(func).is_asynchronous
+        wrapped = asynchronous if is_async else sync
+        cached = (
+            func_in_tracing_context_async if is_async else func_in_tracing_context_sync
+        )
+        for attribute in (
+            "cache_info",
+            "cache_clear",
+            "cache_parameters",
+            "invalidate",
+            "close",
+            "open",
+        ):
+            if hasattr(cached, attribute):
+                setattr(wrapped, attribute, getattr(cached, attribute))
+        return cast(Callable[..., T], wrapped)
+
     def _bootstrap_rest_api(
         self,
     ) -> None:

great_ai/helper/frozen_call.py

@@ -0,0 +1,40 @@
+from copy import deepcopy
+from typing import Any, Dict, Tuple
+
+from pydantic import BaseModel
+
+
+def make_hashable(value: Any) -> Any:
+    """Build an immutable value key without retaining mutable container contents."""
+    if isinstance(value, BaseModel):
+        return type(value), make_hashable(value.model_dump())
+    if isinstance(value, dict):
+        return dict, frozenset(
+            (make_hashable(k), make_hashable(v)) for k, v in value.items()
+        )
+    if isinstance(value, (list, tuple)):
+        return (list if isinstance(value, list) else tuple), tuple(
+            make_hashable(v) for v in value
+        )
+    if isinstance(value, (set, frozenset)):
+        return set, frozenset(make_hashable(v) for v in value)
+    return type(value), value
+
+
+class FrozenCall:
+    """Keep immutable cache identity separate from arguments used by a prediction."""
+
+    def __init__(
+        self, args: Tuple[Any, ...], kwargs: Dict[str, Any], do_not_persist: bool
+    ) -> None:
+        self._key = make_hashable((args, kwargs, do_not_persist))
+        self.args, self.kwargs = deepcopy((args, kwargs))
+        self.do_not_persist = do_not_persist
+
+    def __hash__(self) -> int:
+        return hash(self._key)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, FrozenCall):
+            return NotImplemented
+        return bool(self._key == other._key)

tests/test_cache_regressions.py

@@ -1,8 +1,7 @@
 import asyncio
-from typing import Dict, Tuple
+from typing import Dict, List, Tuple
 
 import pytest
-
 from great_ai import GreatAI
 
 
@@ -42,3 +41,35 @@ async def test_async_predictions_retry_after_failure(
     assert (await predict(7)).output == 7
     assert (await predict(7)).output == 7
     assert calls == 2
+
+
+def test_prediction_mutation_does_not_change_cache_keys(isolated_context: None) -> None:
+    calls = 0
+
+    @GreatAI.create
+    def predict(values: List[int]) -> int:
+        nonlocal calls
+        calls += 1
+        values.append(1)
+        return len(values)
+
+    original: List[int] = []
+    assert predict(original).output == 1
+    assert predict([]).output == 1
+    assert original == []
+    assert calls == 1
+
+
+def test_caller_mutation_does_not_change_cached_trace_values(
+    isolated_context: None,
+) -> None:
+    @GreatAI.create
+    def predict(values: List[List[int]]) -> int:
+        return values[0][0]
+
+    original = [[1]]
+    first = predict(original)
+    original[0][0] = 2
+    assert predict(original).output == 2
+    assert predict([[1]]).trace_id == first.trace_id
+    assert first.logged_values["arg:values:value"] == [[1]]