From a79c65757fec5ec3f722e5e6b138f2f14e22fd47 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 27 Oct 2026 20:09:27 +0000 Subject: [PATCH] Separate immutable prediction cache keys from mutable call arguments --- great_ai/deploy/great_ai.py | 59 +++++++++++++++++++++------------ great_ai/helper/frozen_call.py | 40 ++++++++++++++++++++++ tests/test_cache_regressions.py | 35 +++++++++++++++++-- 3 files changed, 111 insertions(+), 23 deletions(-) create mode 100644 great_ai/helper/frozen_call.py diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index 8057b65..107ddac 100644 --- a/great_ai/deploy/great_ai.py +++ b/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,37 +259,53 @@ 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, diff --git a/great_ai/helper/frozen_call.py b/great_ai/helper/frozen_call.py new file mode 100644 index 0000000..2f13e2f --- /dev/null +++ b/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) diff --git a/tests/test_cache_regressions.py b/tests/test_cache_regressions.py index 004dbbe..a84573e 100644 --- a/tests/test_cache_regressions.py +++ b/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]]