litellm/tests/unit/caching/test_valkey_semantic_cache.py

import hashlib
import os
import struct
import subprocess
import sys
import textwrap
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest


from litellm.caching.valkey_semantic_cache import ValkeySemanticCache

_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))


def _make_cache(sync_client=None, async_client=None, similarity_threshold=0.8):
    return ValkeySemanticCache(
        similarity_threshold=similarity_threshold,
        index_name="test_index",
        sync_client=sync_client or MagicMock(),
        async_client=async_client or AsyncMock(),
    )


def _search_result(distance, response='{"content": "Paris"}'):
    return SimpleNamespace(
        docs=[SimpleNamespace(response=response, vector_distance=str(distance))]
    )


def test_build_valkey_url_prefers_valkey_env(monkeypatch):
    monkeypatch.setenv("REDIS_HOST", "redis-host")
    monkeypatch.setenv("REDIS_PORT", "6379")
    monkeypatch.setenv("REDIS_PASSWORD", "rpass")
    monkeypatch.setenv("VALKEY_HOST", "valkey-host")
    monkeypatch.setenv("VALKEY_PORT", "6380")
    monkeypatch.setenv("VALKEY_PASSWORD", "vpass")

    assert (
        ValkeySemanticCache._build_valkey_url(None, None, None)
        == "redis://:vpass@valkey-host:6380"
    )


def test_build_valkey_url_supports_passwordless(monkeypatch):
    monkeypatch.delenv("REDIS_PASSWORD", raising=False)
    monkeypatch.delenv("VALKEY_PASSWORD", raising=False)
    monkeypatch.setenv("VALKEY_HOST", "valkey-host")
    monkeypatch.setenv("VALKEY_PORT", "6380")

    assert (
        ValkeySemanticCache._build_valkey_url(None, None, None)
        == "redis://valkey-host:6380"
    )


def test_build_valkey_url_falls_back_to_redis_env(monkeypatch):
    monkeypatch.delenv("VALKEY_HOST", raising=False)
    monkeypatch.delenv("VALKEY_PORT", raising=False)
    monkeypatch.delenv("VALKEY_PASSWORD", raising=False)
    monkeypatch.setenv("REDIS_HOST", "redis-host")
    monkeypatch.setenv("REDIS_PORT", "6379")
    monkeypatch.setenv("REDIS_PASSWORD", "rpass")

    assert (
        ValkeySemanticCache._build_valkey_url(None, None, None)
        == "redis://:rpass@redis-host:6379"
    )


def test_build_valkey_url_requires_host_and_port(monkeypatch):
    for var in (
        "VALKEY_HOST",
        "VALKEY_PORT",
        "VALKEY_PASSWORD",
        "REDIS_HOST",
        "REDIS_PORT",
        "REDIS_PASSWORD",
    ):
        monkeypatch.delenv(var, raising=False)

    with pytest.raises(ValueError, match="Missing required Valkey configuration"):
        ValkeySemanticCache._build_valkey_url(None, None, None)


def test_build_valkey_url_uses_rediss_scheme_when_ssl(monkeypatch):
    monkeypatch.setenv("VALKEY_HOST", "valkey-host")
    monkeypatch.setenv("VALKEY_PORT", "6379")
    monkeypatch.setenv("VALKEY_PASSWORD", "vpass")

    assert (
        ValkeySemanticCache._build_valkey_url(None, None, None, ssl=True)
        == "rediss://:vpass@valkey-host:6379"
    )
    assert ValkeySemanticCache._build_valkey_url(
        "h", "6379", None, ssl=False
    ).startswith("redis://")


def test_init_requires_similarity_threshold():
    with pytest.raises(ValueError, match="similarity_threshold must be provided"):
        ValkeySemanticCache(sync_client=MagicMock(), async_client=AsyncMock())


def test_init_stores_embedding_max_input_tokens():
    cache = ValkeySemanticCache(
        similarity_threshold=0.8,
        sync_client=MagicMock(),
        async_client=AsyncMock(),
        embedding_max_input_tokens=512,
    )
    assert cache.embedding_max_input_tokens == 512
    assert _make_cache().embedding_max_input_tokens is None


def test_init_rejects_cluster_startup_nodes():
    with pytest.raises(ValueError, match="cluster-mode-enabled"):
        ValkeySemanticCache(
            similarity_threshold=0.8,
            startup_nodes=[{"host": "shard1", "port": 6379}],
        )


def test_cache_dispatch_rejects_cluster_for_valkey_semantic():
    from litellm.caching.caching import Cache
    from litellm.types.caching import LiteLLMCacheType

    with pytest.raises(ValueError, match="cluster-mode-enabled"):
        Cache(
            type=LiteLLMCacheType.VALKEY_SEMANTIC,
            host="valkey-host",
            port="6379",
            similarity_threshold=0.8,
            redis_startup_nodes=[{"host": "shard1", "port": 6379}],
        )


def test_scope_tag_is_deterministic_hex():
    tag = ValkeySemanticCache._scope_tag("model:gpt-4o::abc-123")
    assert tag == hashlib.sha256(b"model:gpt-4o::abc-123").hexdigest()
    assert len(tag) == 64
    assert ValkeySemanticCache._scope_tag("a") != ValkeySemanticCache._scope_tag("b")


def test_embedding_to_bytes_is_little_endian_float32():
    assert ValkeySemanticCache._embedding_to_bytes([1.0, 0.0]) == struct.pack(
        "<2f", 1.0, 0.0
    )


def test_set_cache_stores_scoped_doc_with_embedding(monkeypatch):
    sync_client = MagicMock()
    cache = _make_cache(sync_client=sync_client)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    cache.set_cache(
        key="cache-key",
        value={"content": "Paris"},
        messages=[{"role": "user", "content": "What is the capital of France?"}],
    )

    sync_client.ft.return_value.create_index.assert_called_once()
    assert sync_client.hset.call_count == 1
    doc_key, kwargs = (
        sync_client.hset.call_args.args[0],
        sync_client.hset.call_args.kwargs,
    )
    mapping = kwargs["mapping"]
    scope = ValkeySemanticCache._scope_tag("cache-key")
    assert mapping[ValkeySemanticCache.CACHE_KEY_FIELD_NAME] == scope
    assert mapping["prompt"] == "What is the capital of France?"
    assert mapping["response"] == "{'content': 'Paris'}"
    assert mapping["embedding"] == struct.pack("<3f", 0.1, 0.2, 0.3)
    assert doc_key.startswith(f"test_index:{scope}:")


def test_set_cache_applies_ttl():
    sync_client = MagicMock()
    cache = _make_cache(sync_client=sync_client)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    cache.set_cache(
        key="cache-key",
        value={"content": "Paris"},
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        ttl=60,
    )

    sync_client.expire.assert_called_once()
    assert sync_client.expire.call_args.args[1] == 60


def test_set_cache_skips_ttl_when_absent():
    sync_client = MagicMock()
    cache = _make_cache(sync_client=sync_client)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    cache.set_cache(
        key="cache-key",
        value={"content": "Paris"},
        messages=[{"role": "user", "content": "What is the capital of France?"}],
    )

    sync_client.expire.assert_not_called()


def test_get_cache_returns_hit_above_threshold():
    sync_client = MagicMock()
    sync_client.ft.return_value.search.return_value = _search_result(0.1)
    cache = _make_cache(sync_client=sync_client, similarity_threshold=0.8)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    metadata = {}
    result = cache.get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital of France?"}],
        metadata=metadata,
    )

    assert result == {"content": "Paris"}
    assert metadata["semantic-similarity"] == pytest.approx(0.9)


def test_get_cache_misses_below_threshold():
    sync_client = MagicMock()
    sync_client.ft.return_value.search.return_value = _search_result(0.5)
    cache = _make_cache(sync_client=sync_client, similarity_threshold=0.8)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    metadata = {}
    result = cache.get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital of Germany?"}],
        metadata=metadata,
    )

    assert result is None
    assert metadata["semantic-similarity"] == pytest.approx(0.5)


def test_get_cache_misses_when_no_docs():
    sync_client = MagicMock()
    sync_client.ft.return_value.search.return_value = SimpleNamespace(docs=[])
    cache = _make_cache(sync_client=sync_client)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    metadata = {}
    result = cache.get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital of France?"}],
        metadata=metadata,
    )

    assert result is None
    assert metadata["semantic-similarity"] == 0.0


def test_get_cache_query_filters_by_scope_tag():
    sync_client = MagicMock()
    sync_client.ft.return_value.search.return_value = _search_result(0.1)
    cache = _make_cache(sync_client=sync_client)
    cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])

    cache.get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital of France?"}],
        metadata={},
    )

    query = sync_client.ft.return_value.search.call_args.args[0]
    scope = ValkeySemanticCache._scope_tag("cache-key")
    assert scope in query.query_string()
    assert "KNN 1 @embedding" in query.query_string()


def _async_ft(search_distance):
    search_obj = SimpleNamespace(
        search=AsyncMock(return_value=_search_result(search_distance)),
        create_index=AsyncMock(),
    )
    return MagicMock(return_value=search_obj)


@pytest.mark.asyncio
async def test_async_set_and_get_roundtrip():
    async_client = AsyncMock()
    async_client.ft = _async_ft(0.05)
    cache = _make_cache(async_client=async_client, similarity_threshold=0.8)
    cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])

    await cache.async_set_cache(
        key="cache-key",
        value={"content": "Paris"},
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        ttl=30,
    )
    async_client.hset.assert_awaited_once()
    async_client.expire.assert_awaited_once()
    assert async_client.expire.call_args.args[1] == 30

    metadata = {}
    result = await cache.async_get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital city of France"}],
        metadata=metadata,
    )
    assert result == {"content": "Paris"}
    assert metadata["semantic-similarity"] == pytest.approx(0.95)


@pytest.mark.asyncio
async def test_async_set_cache_passes_only_metadata_to_get_async_embedding():
    async_client = AsyncMock()
    async_client.ft = _async_ft(0.05)
    cache = _make_cache(async_client=async_client)
    captured: dict[str, object] = {}

    async def spy_embedding(prompt: str, metadata: dict | None = None) -> list[float]:
        captured["prompt"] = prompt
        captured["metadata"] = metadata
        return [0.1, 0.2, 0.3]

    cache._get_async_embedding = spy_embedding

    await cache.async_set_cache(
        key="cache-key",
        value={"content": "Paris"},
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        metadata={"user_api_key": "sk-test"},
        cache_key="abc123",
        custom_llm_provider="openai",
    )

    assert captured["metadata"] == {"user_api_key": "sk-test"}
    async_client.hset.assert_awaited_once()


@pytest.mark.asyncio
async def test_async_get_cache_passes_only_metadata_to_get_async_embedding():
    async_client = AsyncMock()
    async_client.ft = _async_ft(0.05)
    cache = _make_cache(async_client=async_client)
    captured: dict[str, object] = {}

    async def spy_embedding(prompt: str, metadata: dict | None = None) -> list[float]:
        captured["prompt"] = prompt
        captured["metadata"] = dict(metadata) if metadata is not None else None
        return [0.1, 0.2, 0.3]

    cache._get_async_embedding = spy_embedding

    result = await cache.async_get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        metadata={"user_api_key": "sk-test"},
        cache_key="abc123",
        custom_llm_provider="openai",
    )

    assert result == {"content": "Paris"}
    assert captured["metadata"] == {"user_api_key": "sk-test"}


@pytest.mark.asyncio
async def test_async_get_cache_misses_below_threshold():
    async_client = AsyncMock()
    async_client.ft = _async_ft(0.4)
    cache = _make_cache(async_client=async_client, similarity_threshold=0.8)
    cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])

    metadata = {}
    result = await cache.async_get_cache(
        key="cache-key",
        messages=[{"role": "user", "content": "capital of Germany?"}],
        metadata=metadata,
    )
    assert result is None
    assert metadata["semantic-similarity"] == pytest.approx(0.6)


def test_ensure_index_swallows_already_exists():
    sync_client = MagicMock()
    sync_client.ft.return_value.create_index.side_effect = Exception(
        "Index test_index already exists."
    )
    cache = _make_cache(sync_client=sync_client)

    cache._ensure_index_sync(3)
    assert cache._index_dim == 3


def test_ensure_index_reraises_unexpected_error():
    sync_client = MagicMock()
    sync_client.ft.return_value.create_index.side_effect = Exception(
        "connection refused"
    )
    cache = _make_cache(sync_client=sync_client)

    with pytest.raises(Exception, match="connection refused"):
        cache._ensure_index_sync(3)


_FT_INFO_ATTRS_DIM_1536 = [
    [b"identifier", b"litellm_cache_key", b"type", b"TAG"],
    [
        b"identifier",
        b"embedding",
        b"type",
        b"VECTOR",
        b"index",
        [b"capacity", 10240, b"dimensions", 1536, b"distance_metric", b"COSINE"],
    ],
]


def test_extract_index_dim_parses_nested_ft_info():
    info = {"attributes": _FT_INFO_ATTRS_DIM_1536}
    assert ValkeySemanticCache._extract_index_dim(info) == 1536
    assert ValkeySemanticCache._extract_index_dim({"attributes": []}) is None


def test_ensure_index_raises_on_dimension_mismatch():
    sync_client = MagicMock()
    sync_client.ft.return_value.create_index.side_effect = Exception(
        "Index test_index already exists."
    )
    sync_client.ft.return_value.info.return_value = {
        "attributes": _FT_INFO_ATTRS_DIM_1536
    }
    cache = _make_cache(sync_client=sync_client)

    with pytest.raises(
        ValueError, match="already exists with embedding dimension 1536"
    ):
        cache._ensure_index_sync(768)
    assert cache._index_dim is None


def test_ensure_index_accepts_matching_existing_dimension():
    sync_client = MagicMock()
    sync_client.ft.return_value.create_index.side_effect = Exception(
        "Index test_index already exists."
    )
    sync_client.ft.return_value.info.return_value = {
        "attributes": _FT_INFO_ATTRS_DIM_1536
    }
    cache = _make_cache(sync_client=sync_client)

    cache._ensure_index_sync(1536)
    assert cache._index_dim == 1536


def test_init_builds_only_missing_client_from_url():
    sync_client = MagicMock()
    cache = ValkeySemanticCache(
        similarity_threshold=0.8,
        redis_url="redis://valkey-host:6380",
        sync_client=sync_client,
    )
    assert cache.sync_client is sync_client
    assert cache.async_client is not None and cache.async_client is not sync_client


def test_init_uses_both_injected_clients_without_connection_info(monkeypatch):
    for var in ("VALKEY_HOST", "VALKEY_PORT", "REDIS_HOST", "REDIS_PORT"):
        monkeypatch.delenv(var, raising=False)
    sync_client = MagicMock()
    async_client = AsyncMock()

    cache = ValkeySemanticCache(
        similarity_threshold=0.8,
        sync_client=sync_client,
        async_client=async_client,
    )

    assert cache.sync_client is sync_client
    assert cache.async_client is async_client


def test_cache_dispatches_valkey_semantic_type():
    from litellm.caching.caching import Cache
    from litellm.types.caching import LiteLLMCacheType

    cache = Cache(
        type=LiteLLMCacheType.VALKEY_SEMANTIC,
        host="valkey-host",
        port="6380",
        similarity_threshold=0.8,
    )

    assert isinstance(cache.cache, ValkeySemanticCache)


@pytest.mark.asyncio
async def test_index_info_uses_valkey_ft_info():
    # The /health/readiness endpoint calls _index_info() on any
    # RedisSemanticCache instance; since ValkeySemanticCache subclasses it,
    # the inherited RedisVL implementation (which reads self.llmcache) would
    # break. This override must query valkey-search FT.INFO instead.
    async_client = AsyncMock()
    info_namespace = SimpleNamespace(info=AsyncMock(return_value={"num_docs": 3}))
    async_client.ft = MagicMock(return_value=info_namespace)
    cache = _make_cache(async_client=async_client)

    result = await cache._index_info()

    assert result == {"num_docs": 3}
    async_client.ft.assert_called_once_with("test_index")


def test_importing_caching_does_not_require_redis():
    # redis is an optional dependency (extra_proxy), so the base SDK can be
    # installed without it. Selecting valkey-semantic needs redis, but merely
    # importing litellm.caching.caching must not, or `import litellm` breaks for
    # every base-SDK user. This runs in a subprocess with redis blocked so the
    # check is not polluted by redis already being imported in this session.
    code = textwrap.dedent("""
        import sys
        for name in ("redis", "redis.asyncio", "redis.commands",
                     "redis.commands.search"):
            sys.modules[name] = None
        import litellm.caching.caching  # must not import redis at module top
        from litellm.types.caching import LiteLLMCacheType
        assert LiteLLMCacheType.VALKEY_SEMANTIC == "valkey-semantic"
        print("ok")
        """)
    result = subprocess.run(
        [sys.executable, "-c", code],
        capture_output=True,
        text=True,
        env={**os.environ, "PYTHONPATH": _REPO_ROOT},
    )
    assert result.returncode == 0, result.stderr
    assert "ok" in result.stdout