import asyncio from collections.abc import Callable from pathlib import Path from typing import Final, TypeAlias, cast from urllib.parse import urlparse from uuid import uuid4 import pytest from litellm.caching.caching import Cache from litellm.rust_bridge import _native from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime from litellm.types.caching import LiteLLMCacheType from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.cache import activate_native, assert_native_runtime, completion_kwargs from tests.test_litellm_rust.support.s3_stub import S3Stub pytestmark: Final = pytest.mark.requires_rust_extension CacheFactory: TypeAlias = Callable[[], Cache] @pytest.fixture def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: backend: Final = cast(LiteLLMCacheType, request.param) match backend: case LiteLLMCacheType.LOCAL: return lambda: Cache(type=backend) case LiteLLMCacheType.DISK: return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) case LiteLLMCacheType.REDIS: parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) case LiteLLMCacheType.S3: stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) return lambda: Cache( type=backend, s3_bucket_name="cache-bucket", s3_region_name="us-east-1", s3_endpoint_url=stub.url, s3_aws_access_key_id="key", s3_aws_secret_access_key="secret", s3_path="team", ) case LiteLLMCacheType.GCS: return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") case LiteLLMCacheType.REDIS_SEMANTIC: return lambda: Cache( type=backend, redis_url="redis://127.0.0.1:6379", similarity_threshold=0.8, redis_semantic_cache_embedding_model="text-embedding-3-small", ) case LiteLLMCacheType.VALKEY_SEMANTIC: return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) case _: raise AssertionError(f"no local factory for {backend}") ROUND_TRIP_BACKENDS: Final = ( LiteLLMCacheType.LOCAL, LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3, ) SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) @pytest.mark.parametrize( "cache_factory", [ LiteLLMCacheType.LOCAL, LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3, LiteLLMCacheType.GCS, LiteLLMCacheType.REDIS_SEMANTIC, LiteLLMCacheType.VALKEY_SEMANTIC, ], indirect=True, ) def test_legacy_constructor_keeps_python_backends(cache_factory: CacheFactory) -> None: assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor @pytest.mark.parametrize( "cache_factory", [ LiteLLMCacheType.LOCAL, LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3, LiteLLMCacheType.GCS, LiteLLMCacheType.REDIS_SEMANTIC, LiteLLMCacheType.VALKEY_SEMANTIC, ], indirect=True, ) def test_explicit_selection_activates_the_native_backend( cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest ) -> None: assert_native_runtime(activate_native(cache_factory())) @pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) async def test_facade_storage_calls_round_trip_through_the_native_backend( cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest ) -> None: facade: Final = activate_native(cache_factory()) assert_native_runtime(facade) sync_kwargs: Final = completion_kwargs("sync") facade.add_cache({"answer": 1}, **sync_kwargs) assert facade.get_cache(**sync_kwargs) == {"answer": 1} async_kwargs: Final = completion_kwargs("async") await facade.async_add_cache({"answer": 2}, **async_kwargs) assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} assert facade.get_cache(**completion_kwargs("absent")) is None async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: facade: Final = activate_native(Cache(type=LiteLLMCacheType.LOCAL)) assert_native_runtime(facade) kwargs: Final = completion_kwargs("memory") facade.add_cache({"answer": 1}, **kwargs) assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None assert facade.get_cache(**kwargs) == {"answer": 1} @pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) async def test_native_and_python_facades_share_one_wire_format( cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest ) -> None: python_facade: Final = cache_factory() assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor native_facade: Final = activate_native(cache_factory()) assert_native_runtime(native_facade) native_written: Final = completion_kwargs("native") native_facade.add_cache({"writer": "native"}, **native_written) assert python_facade.get_cache(**native_written) == {"writer": "native"} python_written: Final = completion_kwargs("python") python_facade.add_cache({"writer": "python"}, **python_written) assert native_facade.get_cache(**python_written) == {"writer": "python"} async_native: Final = completion_kwargs("async-native") await native_facade.async_add_cache({"writer": "async-native"}, **async_native) assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} async_python: Final = completion_kwargs("async-python") await python_facade.async_add_cache({"writer": "async-python"}, **async_python) assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} @pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) async def test_embedding_pipeline_stores_one_native_entry_per_input( cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest ) -> None: facade: Final = activate_native(cache_factory()) assert_native_runtime(facade) inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] result: Final = EmbeddingResponse( model="text-embedding-3-small", data=[ {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, ], ) await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] assert len(set(keys)) == len(inputs) for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) assert isinstance(cached, dict) assert cached["embedding"] == expected assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None @pytest.mark.parametrize( ("backend", "settings", "message"), [ pytest.param( LiteLLMCacheType.VALKEY_SEMANTIC, {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, "native Valkey semantic cache does not support TLS connections", id="valkey-tls", ), pytest.param( LiteLLMCacheType.VALKEY_SEMANTIC, {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", id="valkey-socket-timeout", ), pytest.param( LiteLLMCacheType.REDIS_SEMANTIC, {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, "native Redis semantic cache does not support TLS or query options in redis_url", id="redis-semantic-tls", ), pytest.param( LiteLLMCacheType.REDIS_SEMANTIC, {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, "native Redis semantic cache does not support TLS or query options in redis_url", id="redis-semantic-query", ), ], ) def test_semantic_settings_the_native_client_cannot_honor_decline( monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str ) -> None: with pytest.raises(_native.RustBridgeDeclined, match=message): activate_native(Cache(type=backend, **settings)) class _SemanticHit: """A native semantic runtime that answers every lookup with one cached response.""" kind: Final = "native" def lookup_semantic(self, request: object) -> tuple[object, float | None]: return {"answer": 42}, 0.97 async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: return {"answer": 42}, 0.97 @pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) @pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) def test_native_semantic_hit_stamps_similarity_on_request_metadata( semantic_type: LiteLLMCacheType, use_async: bool ) -> None: """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the facade copies it to the caller's metadata; the native path must report it the same way.""" facade: Final = Cache() facade.type = semantic_type facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter metadata: Final[dict[str, object]] = {} kwargs: Final = { "cache_key": "semantic-key", "messages": [{"role": "user", "content": "hello"}], "metadata": metadata, } result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) assert result == {"answer": 42} assert metadata["semantic-similarity"] == 0.97