refactor(anthropic): move prompt cache observation lookup out of the proxy hook (#45748)

* refactor(anthropic): move prompt cache observation lookup out of the proxy hook

Part of LIT-9423

* test(internal_context): exempt moved prompt cache lookup helper from key-family check

litellm/llms/anthropic/cache_aware_routing.py

@@ -19,12 +19,12 @@ from litellm.llms.anthropic.prompt_cache_prediction import (
     UnsupportedPredictionTarget,
     cache_scope,
     count_prompt_tokens,
+    lookup,
     parse_prompt,
     resolve_prediction_target,
     supported_prediction_headers,
 )
 from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
-from litellm.proxy.hooks.prompt_cache_prediction import lookup
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.management_endpoints.prompt_cache_prediction import (
     CacheCostScenario,

litellm/llms/anthropic/prompt_cache_prediction.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import hashlib
 import json
+import time
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from functools import reduce
@@ -14,6 +15,7 @@ import httpx
 from pydantic import ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
 
 import litellm
+from litellm.caching.dual_cache import DualCache
 from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
 from litellm.llms.anthropic.common_utils import (
     AnthropicModelInfo,
@@ -771,3 +773,50 @@ def parse_observed_cache(
         cached_tokens=cache_tokens,
         cache_creation_tokens=usage.cache_creation_tokens,
     )
+
+
+_RETENTION_SECONDS: Final = 86_400
+
+
+class CacheObservation(LiteLLMBaseModel):
+    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
+
+    fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
+    cached_tokens: int = Field(gt=0)
+    observed_at: float = Field(ge=0, allow_inf_nan=False)
+    expires_at: float = Field(ge=0, allow_inf_nan=False)
+
+
+_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
+
+
+def _cache_key(scope: str, fingerprint: str) -> str:
+    return f"prompt-cache-observation:{scope}:{fingerprint}"
+
+
+async def lookup(
+    cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
+) -> CacheObservation | None:
+    checked_at: Final = time.time() if now is None else now
+    exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
+    if exact is not None and exact.expires_at > checked_at:
+        return exact
+    older: Final = await asyncio.gather(
+        *(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
+    )
+    observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
+    return next(
+        (observation for observation in observations if observation.expires_at > checked_at),
+        next(iter(observations), None),
+    )
+
+
+async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
+    try:
+        value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # validate the legacy cache's untyped result at the I/O boundary
+        if value is None:
+            return None
+        observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
+    except ValidationError:
+        return None
+    return observation if observation.fingerprint == fingerprint else None

litellm/proxy/hooks/prompt_cache_prediction.py

@@ -1,70 +1,28 @@
 from __future__ import annotations
 
-import asyncio
 import time
 from collections.abc import Callable, Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, Final, Literal
 
 import httpx
-from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
+from pydantic import ConfigDict, Field, ValidationError
 
 from litellm._internal_context import with_service_target
-from litellm.caching.dual_cache import DualCache
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, parse_observed_cache
+from litellm.llms.anthropic.prompt_cache_prediction import (
+    _RETENTION_SECONDS,  # pyright: ignore[reportPrivateUsage]  # moved to the SDK layer, still shared with this hook
+    CacheObservation,
+    _cache_key,  # pyright: ignore[reportPrivateUsage]  # moved to the SDK layer, still shared with this hook
+    _read_exact,  # pyright: ignore[reportPrivateUsage]  # moved to the SDK layer, still shared with this hook
+    parse_observed_cache,
+)
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.utils import ModelResponse
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import InternalUsageCache
 
-_RETENTION_SECONDS: Final = 86_400
-
-
-class CacheObservation(LiteLLMBaseModel):
-    model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
-
-    fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
-    cached_tokens: int = Field(gt=0)
-    observed_at: float = Field(ge=0, allow_inf_nan=False)
-    expires_at: float = Field(ge=0, allow_inf_nan=False)
-
-
-_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
-
-
-def _cache_key(scope: str, fingerprint: str) -> str:
-    return f"prompt-cache-observation:{scope}:{fingerprint}"
-
-
-async def lookup(
-    cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
-) -> CacheObservation | None:
-    checked_at: Final = time.time() if now is None else now
-    exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
-    if exact is not None and exact.expires_at > checked_at:
-        return exact
-    older: Final = await asyncio.gather(
-        *(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
-    )
-    observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
-    return next(
-        (observation for observation in observations if observation.expires_at > checked_at),
-        next(iter(observations), None),
-    )
-
-
-async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
-    try:
-        value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # validate the legacy cache's untyped result at the I/O boundary
-        if value is None:
-            return None
-        observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
-    except ValidationError:
-        return None
-    return observation if observation.fingerprint == fingerprint else None
-
 
 class _Metadata(LiteLLMBaseModel):
     model_config = ConfigDict(strict=True)

tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py

@@ -22,6 +22,7 @@ from litellm.llms.anthropic.prompt_cache_prediction import (
     cache_scope,
     count_cache_plan,
     count_prompt_tokens,
+    lookup,
     parse_cache_plan,
     parse_observed_cache,
     parse_prompt,
@@ -32,7 +33,7 @@ from litellm.llms.anthropic.prompt_cache_prediction import (
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.models.credentials import CredentialItem
 from litellm.proxy import proxy_server
-from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver, lookup
+from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver
 from litellm.proxy.management_endpoints.prompt_cache_prediction import predict_arm
 from litellm.proxy.utils import InternalUsageCache
 from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo

tests/unit/proxy/hooks/test_prompt_cache_observer.py

@@ -10,11 +10,8 @@ import litellm
 from litellm.caching.dual_cache import DualCache
 from litellm.llms.anthropic.chat.transformation import AnthropicConfig
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
-from litellm.llms.anthropic.prompt_cache_prediction import cache_scope, parse_prompt
-from litellm.proxy.hooks.prompt_cache_prediction import (
-    PromptCacheObserver,
-    lookup,
-)
+from litellm.llms.anthropic.prompt_cache_prediction import cache_scope, lookup, parse_prompt
+from litellm.proxy.hooks.prompt_cache_prediction import PromptCacheObserver
 from litellm.proxy.utils import InternalUsageCache
 from litellm.types.utils import ModelResponse
 

tests/unit/test_internal_context.py

@@ -29,6 +29,7 @@ _CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache")
 # Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade.
 _CACHE_PARAMETER_HELPERS: Final = frozenset(
     {
+        "litellm/llms/anthropic/prompt_cache_prediction.py",
         "litellm/proxy/common_utils/cache_coordinator.py",
         "litellm/proxy/common_utils/user_api_key_cache.py",
         "litellm/utils.py",