From ebf2156ec4a64996fbf03d27412d9fdd116fd783 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 12:06:16 -0700 Subject: [PATCH] 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 | 2 +- .../llms/anthropic/prompt_cache_prediction.py | 49 ++++++++++++++++ .../proxy/hooks/prompt_cache_prediction.py | 58 +++---------------- .../test_anthropic_prompt_cache_prediction.py | 3 +- .../proxy/hooks/test_prompt_cache_observer.py | 7 +-- tests/unit/test_internal_context.py | 1 + 6 files changed, 63 insertions(+), 57 deletions(-) diff --git a/litellm/llms/anthropic/cache_aware_routing.py b/litellm/llms/anthropic/cache_aware_routing.py index 024631c961a..f049e4ac187 100644 --- a/litellm/llms/anthropic/cache_aware_routing.py +++ b/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, diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 634eb264b38..4c363dba860 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/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 diff --git a/litellm/proxy/hooks/prompt_cache_prediction.py b/litellm/proxy/hooks/prompt_cache_prediction.py index b4f7322b58b..ec1cd421af7 100644 --- a/litellm/proxy/hooks/prompt_cache_prediction.py +++ b/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) diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index 9bb26b66aa7..df89774951f 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/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 diff --git a/tests/unit/proxy/hooks/test_prompt_cache_observer.py b/tests/unit/proxy/hooks/test_prompt_cache_observer.py index 82af3e9a6ef..3354aa457b7 100644 --- a/tests/unit/proxy/hooks/test_prompt_cache_observer.py +++ b/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 diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index ccbe1487fc4..200d9a0000c 100644 --- a/tests/unit/test_internal_context.py +++ b/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",