litellm/tests/unit/caching/test_unit_test_caching.py

import time
import traceback
from litellm._uuid import uuid

from dotenv import load_dotenv

load_dotenv()
import asyncio
import hashlib
import random
from typing import Final

import pytest

import litellm
from litellm import aembedding, completion, embedding
from litellm.caching.caching import Cache

from unittest.mock import AsyncMock, patch, MagicMock
from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse, _PENDING_CACHE_WRITES
from litellm.caching.caching import LiteLLMCacheType
from litellm.types.utils import CallTypes
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
    ModelResponse,
    EmbeddingResponse,
    TextCompletionResponse,
    TranscriptionResponse,
    Embedding,
)
from datetime import timedelta, datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
from litellm._logging import verbose_logger
import logging


def test_get_kwargs_for_cache_key():
    _cache = litellm.Cache()
    relevant_kwargs = ModelParamHelper.get_all_llm_api_params()
    print(relevant_kwargs)


def test_get_cache_key_chat_completion():
    cache = Cache()
    kwargs = {
        "model": "gpt-3.5-turbo",
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "temperature": 0.7,
    }
    cache_key_1 = cache.get_cache_key(**kwargs)
    assert isinstance(cache_key_1, str)
    assert len(cache_key_1) > 0

    kwargs_2 = {
        "model": "gpt-3.5-turbo",
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "max_completion_tokens": 100,
    }
    cache_key_2 = cache.get_cache_key(**kwargs_2)
    assert cache_key_1 != cache_key_2

    kwargs_3 = {
        "model": "gpt-3.5-turbo",
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "max_completion_tokens": 100,
    }
    cache_key_3 = cache.get_cache_key(**kwargs_3)
    assert cache_key_2 == cache_key_3


def test_get_cache_key_embedding():
    cache = Cache()
    kwargs = {
        "model": "text-embedding-3-small",
        "input": "Hello, world!",
        "dimensions": 1536,
    }
    cache_key_1 = cache.get_cache_key(**kwargs)
    assert isinstance(cache_key_1, str)
    assert len(cache_key_1) > 0

    kwargs_2 = {
        "model": "text-embedding-3-small",
        "input": "Hello, world!",
        "dimensions": 1539,
    }
    cache_key_2 = cache.get_cache_key(**kwargs_2)
    assert cache_key_1 != cache_key_2

    kwargs_3 = {
        "model": "text-embedding-3-small",
        "input": "Hello, world!",
        "dimensions": 1539,
    }
    cache_key_3 = cache.get_cache_key(**kwargs_3)
    assert cache_key_2 == cache_key_3


def test_get_cache_key_text_completion():
    cache = Cache()
    kwargs = {
        "model": "gpt-3.5-turbo",
        "prompt": "Hello, world! here is a second line",
        "best_of": 3,
        "logit_bias": {"123": 1},
        "seed": 42,
    }
    cache_key_1 = cache.get_cache_key(**kwargs)
    assert isinstance(cache_key_1, str)
    assert len(cache_key_1) > 0

    kwargs_2 = {
        "model": "gpt-3.5-turbo",
        "prompt": "Hello, world! here is a second line",
        "best_of": 30,
    }
    cache_key_2 = cache.get_cache_key(**kwargs_2)
    assert cache_key_1 != cache_key_2

    kwargs_3 = {
        "model": "gpt-3.5-turbo",
        "prompt": "Hello, world! here is a second line",
        "best_of": 30,
    }
    cache_key_3 = cache.get_cache_key(**kwargs_3)
    assert cache_key_2 == cache_key_3


def test_get_cache_key_responses_api():
    """
    Regression test: two /v1/responses calls that differ only in
    `instructions` (or any Responses-API-only param) must produce
    different cache keys. Mirrors the chat / embedding / text-completion
    cache-key tests above.
    """
    cache = Cache()

    base_kwargs = {
        "model": "openai/gpt-4.1",
        "input": [{"role": "user", "content": "what is the weather"}],
        "temperature": 0.3,
    }

    kwargs_a = {
        **base_kwargs,
        "instructions": "summarize the weather on 10th May",
    }
    kwargs_b = {
        **base_kwargs,
        "instructions": "summarize the weather on 7th May",
    }

    key_a = cache.get_cache_key(**kwargs_a)
    key_b = cache.get_cache_key(**kwargs_b)

    assert isinstance(key_a, str) and len(key_a) > 0
    assert key_a != key_b, "instructions must be part of the Responses API cache key"

    # Sanity: identical payloads must still collide (cache hits still work)
    key_a_again = cache.get_cache_key(**kwargs_a)
    assert key_a == key_a_again

    # Spot-check a handful of other Responses-only params individually.
    for param, value_x, value_y in [
        ("previous_response_id", "resp_aaa", "resp_bbb"),
        ("reasoning", {"effort": "low"}, {"effort": "high"}),
        ("include", ["reasoning.encrypted_content"], []),
        ("max_output_tokens", 100, 500),
        ("background", True, False),
    ]:
        kx = {**base_kwargs, param: value_x}
        ky = {**base_kwargs, param: value_y}
        assert cache.get_cache_key(**kx) != cache.get_cache_key(
            **ky
        ), f"Responses-API param `{param}` is not part of the cache key"


def test_get_hashed_cache_key():
    cache = Cache()
    cache_key = "model:gpt-3.5-turbo,messages:Hello world"
    hashed_key = Cache._get_hashed_cache_key(cache_key)
    assert len(hashed_key) == 64  # SHA-256 produces a 64-character hex string


def test_add_namespace_to_cache_key():
    cache = Cache(namespace="test_namespace")
    hashed_key = "abcdef1234567890"

    # Test with class-level namespace
    result = cache._add_namespace_to_cache_key(hashed_key)
    assert result == "test_namespace:abcdef1234567890"

    # Test with metadata namespace
    kwargs = {"metadata": {"redis_namespace": "custom_namespace"}}
    result = cache._add_namespace_to_cache_key(hashed_key, **kwargs)
    assert result == "custom_namespace:abcdef1234567890"

    # Test with cache control namespace
    kwargs = {"cache": {"namespace": "cache_control_namespace"}}
    result = cache._add_namespace_to_cache_key(hashed_key, **kwargs)
    assert result == "cache_control_namespace:abcdef1234567890"

    kwargs = {"cache": {"namespace": "cache_control_namespace-2"}}
    result = cache._add_namespace_to_cache_key(hashed_key, **kwargs)
    assert result == "cache_control_namespace-2:abcdef1234567890"


def test_get_model_param_value():
    cache = Cache()

    # Test with regular model
    kwargs = {"model": "gpt-3.5-turbo"}
    assert cache._get_model_param_value(kwargs) == "gpt-3.5-turbo"

    # Test with model_group
    kwargs = {"model": "gpt-3.5-turbo", "metadata": {"model_group": "gpt-group"}}
    assert cache._get_model_param_value(kwargs) == "gpt-group"

    # Test with caching_group
    kwargs = {
        "model": "gpt-3.5-turbo",
        "metadata": {
            "model_group": "openai-gpt-3.5-turbo",
            "caching_groups": [("openai-gpt-3.5-turbo", "azure-gpt-3.5-turbo")],
        },
    }
    assert (
        cache._get_model_param_value(kwargs)
        == "('openai-gpt-3.5-turbo', 'azure-gpt-3.5-turbo')"
    )

    kwargs = {
        "model": "gpt-3.5-turbo",
        "metadata": {
            "model_group": "azure-gpt-3.5-turbo",
            "caching_groups": [("openai-gpt-3.5-turbo", "azure-gpt-3.5-turbo")],
        },
    }
    assert (
        cache._get_model_param_value(kwargs)
        == "('openai-gpt-3.5-turbo', 'azure-gpt-3.5-turbo')"
    )

    kwargs = {
        "model": "gpt-3.5-turbo",
        "metadata": {
            "model_group": "not-in-caching-group-gpt-3.5-turbo",
            "caching_groups": [("openai-gpt-3.5-turbo", "azure-gpt-3.5-turbo")],
        },
    }
    assert cache._get_model_param_value(kwargs) == "not-in-caching-group-gpt-3.5-turbo"


def test_get_model_param_value_reads_model_group_from_litellm_metadata():
    cache = Cache()
    request = {
        "model": "openai/gpt-5.6",
        "input": "search this text",
        "tools": [{"type": "web_search_preview", "search_context_size": "medium"}],
    }

    assert cache._get_model_param_value({**request, "litellm_metadata": {"model_group": "group-a"}}) == "group-a"
    assert (
        cache._get_model_param_value({**request, "litellm_params": {"litellm_metadata": {"model_group": "group-a"}}})
        == "group-a"
    )
    assert cache._get_model_param_value(
        {
            **request,
            "litellm_metadata": {
                "model_group": "group-a",
                "caching_groups": [("group-a", "group-b")],
            },
        }
    ) == "('group-a', 'group-b')"
    assert cache.get_cache_key(**request, litellm_metadata={"model_group": "group-a"}) != cache.get_cache_key(
        **request, litellm_metadata={"model_group": "group-b"}
    )


def test_get_model_param_value_prefers_the_router_group_over_a_caller_model_group():
    cache: Final = Cache()
    deployment: Final = {"id": "deployment-1"}
    request: Final = {"model": "openai/gpt-5.6", "input": "search this text", "model_info": deployment}
    router_metadata: Final = {"model_group": "caller", "model_info": deployment}

    assert (
        cache._get_model_param_value(
            {**request, "metadata": {"model_group": "target"}, "litellm_metadata": router_metadata}
        )
        == "caller"
    )
    assert (
        cache._get_model_param_value(
            {
                **request,
                "litellm_params": {"metadata": {"model_group": "target"}, "litellm_metadata": router_metadata},
            }
        )
        == "caller"
    )
    assert (
        cache._get_model_param_value(
            {
                **request,
                "metadata": {"model_group": "target", "model_info": deployment},
                "litellm_metadata": router_metadata,
            }
        )
        == "caller"
    )
    assert (
        cache._get_model_param_value(
            {
                **request,
                "metadata": router_metadata,
                "litellm_metadata": {"model_group": "parent", "model_info": {"id": "parent-deployment"}},
            }
        )
        == "caller"
    )


def test_get_model_param_value_keeps_metadata_first_outside_the_router():
    cache: Final = Cache()

    assert (
        cache._get_model_param_value(
            {
                "model": "openai/gpt-5.6",
                "input": "search this text",
                "metadata": {"model_group": "group-a"},
                "litellm_metadata": {"model_group": "group-b"},
            }
        )
        == "group-a"
    )


@pytest.mark.asyncio
async def test_router_call_keeps_its_own_group_when_forwarded_litellm_metadata_names_another(
    monkeypatch: pytest.MonkeyPatch,
):
    monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
    router: Final = litellm.Router(
        model_list=[
            {
                "model_name": group,
                "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-test", "mock_response": f"{group} answer"},
            }
            for group in ("group-a", "group-b")
        ]
    )
    messages: Final = [{"role": "user", "content": "which group answers this"}]

    await router.acompletion(
        model="group-a",
        messages=messages,
        caching=True,
        litellm_metadata={"model_group": "group-b", "model_info": {"id": "parent-deployment"}},
    )
    await asyncio.gather(*_PENDING_CACHE_WRITES)
    reply: Final = await router.acompletion(model="group-b", messages=messages, caching=True)

    assert reply.choices[0].message.content == "group-b answer"


@pytest.mark.parametrize("planted", [["sales-premium"], [("sales", "sales-premium")]], ids=["string", "tuple"])
def test_get_model_param_value_reads_caching_groups_only_beside_the_router_group(planted: list[str | tuple[str, str]]):
    cache: Final = Cache()
    request: Final = {
        "model": "openai/gpt-5.6",
        "input": "search this text",
        "model_info": {"id": "sales-deployment"},
        "litellm_metadata": {"model_group": "sales", "caching_groups": None, "model_info": {"id": "sales-deployment"}},
        "metadata": {"caching_groups": planted},
    }

    assert cache._get_model_param_value(request) == "sales"


def test_preset_cache_key():
    """
    Test that the preset cache key is used if it is set in kwargs["litellm_params"]
    """
    cache = Cache()
    kwargs = {
        "model": "gpt-3.5-turbo",
        "messages": [{"role": "user", "content": "Hello, world!"}],
        "temperature": 0.7,
        "litellm_params": {"preset_cache_key": "preset-cache-key"},
    }

    assert cache.get_cache_key(**kwargs) == "preset-cache-key"


def test_generate_streaming_content():
    cache = Cache()
    content = "Hello, this is a test message."
    generator = cache.generate_streaming_content(content)

    full_response = ""
    chunk_count = 0

    for chunk in generator:
        chunk_count += 1
        assert "choices" in chunk
        assert len(chunk["choices"]) == 1
        assert "delta" in chunk["choices"][0]
        assert "role" in chunk["choices"][0]["delta"]
        assert chunk["choices"][0]["delta"]["role"] == "assistant"
        assert "content" in chunk["choices"][0]["delta"]

        chunk_content = chunk["choices"][0]["delta"]["content"]
        full_response += chunk_content

        # Check that each chunk is no longer than 5 characters
        assert len(chunk_content) <= 5
    print("full_response from generate_streaming_content", full_response)
    # Check that the full content is reconstructed correctly
    assert full_response == content
    # Check that there were multiple chunks
    assert chunk_count > 1

    print(f"Number of chunks: {chunk_count}")