litellm/tests/test_litellm/caching/test_caching_handler.py

import asyncio
import json
import time
from unittest.mock import MagicMock, patch

import httpx
import pytest
import respx
from fastapi.testclient import TestClient

from datetime import datetime
from unittest.mock import AsyncMock

from litellm.caching.caching_handler import LLMCachingHandler


@pytest.mark.asyncio
async def test_process_async_embedding_cached_response():
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    args = {
        "cached_result": [
            {
                "embedding": [-0.025122925639152527, -0.019487135112285614],
                "index": 0,
                "object": "embedding",
            }
        ]
    }

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=args["cached_result"],
        kwargs={"model": "text-embedding-ada-002", "input": "test"},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="text-embedding-ada-002",
    )

    assert cache_hit

    print(f"response: {response}")
    assert len(response.data) == 1


@pytest.mark.asyncio
async def test_embedding_cache_preserves_prompt_tokens_details():
    """Test that prompt_tokens_details (including image_count) survives a full cache hit."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "amazon.titan-embed-image-v1",
            "prompt_tokens_details": {"image_count": 1},
        }
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="amazon.titan-embed-image-v1",
    )

    assert cache_hit
    assert response.usage is not None
    assert response.usage.prompt_tokens_details is not None
    assert response.usage.prompt_tokens_details.image_count == 1


@pytest.mark.asyncio
async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
    """Test that old cached items without prompt_tokens_details still work."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    # Old-format cached item — no prompt_tokens_details field
    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "text-embedding-ada-002",
        }
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "text-embedding-ada-002", "input": "test"},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="text-embedding-ada-002",
    )

    assert cache_hit
    assert response.usage is not None
    assert response.usage.prompt_tokens_details is None


@pytest.mark.asyncio
async def test_embedding_cache_aggregates_multiple_image_counts():
    """Test that image_count is summed correctly across multiple cached items."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "amazon.titan-embed-image-v1",
            "prompt_tokens_details": {"image_count": 1},
        },
        {
            "embedding": [0.031, 0.042],
            "index": 1,
            "object": "embedding",
            "model": "amazon.titan-embed-image-v1",
            "prompt_tokens_details": {"image_count": 1},
        },
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={
            "model": "amazon.titan-embed-image-v1",
            "input": ["img1", "img2"],
        },
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="amazon.titan-embed-image-v1",
    )

    assert cache_hit
    assert response.usage.prompt_tokens_details is not None
    assert response.usage.prompt_tokens_details.image_count == 2


def test_combine_usage_merges_prompt_tokens_details():
    """Test that combine_usage merges prompt_tokens_details from both Usage objects."""
    from litellm.types.utils import PromptTokensDetailsWrapper, Usage

    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    usage1 = Usage(
        prompt_tokens=10,
        completion_tokens=0,
        total_tokens=10,
        prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
    )
    usage2 = Usage(
        prompt_tokens=20,
        completion_tokens=0,
        total_tokens=20,
        prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
    )

    combined = llm_caching_handler.combine_usage(usage1, usage2)

    assert combined.prompt_tokens == 30
    assert combined.total_tokens == 30
    assert combined.prompt_tokens_details is not None
    assert combined.prompt_tokens_details.image_count == 3


def test_combine_usage_handles_none_details():
    """Test that combine_usage works when one or both sides have null prompt_tokens_details."""
    from litellm.types.utils import PromptTokensDetailsWrapper, Usage

    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    # Both null
    usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
    usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
    combined = llm_caching_handler.combine_usage(usage_a, usage_b)
    assert combined.prompt_tokens_details is None

    # Only first has details
    usage_c = Usage(
        prompt_tokens=10,
        completion_tokens=0,
        total_tokens=10,
        prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
    )
    combined = llm_caching_handler.combine_usage(usage_c, usage_b)
    assert combined.prompt_tokens_details is not None
    assert combined.prompt_tokens_details.image_count == 1

    # Only second has details
    combined = llm_caching_handler.combine_usage(usage_a, usage_c)
    assert combined.prompt_tokens_details is not None
    assert combined.prompt_tokens_details.image_count == 1


def test_is_chat_completion_cached_dict():
    from litellm.caching.caching_handler import _is_chat_completion_cached_dict

    assert _is_chat_completion_cached_dict(
        {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []}
    )
    assert _is_chat_completion_cached_dict(
        {"id": "other", "object": "chat.completion.chunk", "choices": []}
    )
    assert _is_chat_completion_cached_dict(
        {"id": "no-object", "choices": [{"index": 0}]}
    )
    assert not _is_chat_completion_cached_dict(
        {"id": "resp_abc", "object": "response", "output": []}
    )


def _build_logging_obj(call_type: str, stream: bool):
    import uuid as _uuid

    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging

    return LiteLLMLogging(
        litellm_call_id=str(datetime.now()),
        call_type=call_type,
        model="gpt-5.4",
        messages=[],
        function_id=str(_uuid.uuid4()),
        stream=stream,
        start_time=datetime.now(),
    )


def test_convert_cached_aresponses_bridge_chat_completion_stream():
    """openai/responses chat-completions bridge: streaming cache hit replays as chat stream."""
    from litellm import aresponses
    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
    from litellm.types.utils import CallTypes

    caching_handler = LLMCachingHandler(
        original_function=aresponses, request_kwargs={}, start_time=datetime.now()
    )
    cached_result = {
        "id": "chatcmpl-bridge-cache-test",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "gpt-5.4",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "Hi!"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
    }

    result = caching_handler._convert_cached_result_to_model_response(
        cached_result=cached_result,
        call_type=CallTypes.aresponses.value,
        kwargs={
            "model": "gpt-5.4",
            "stream": True,
            "messages": [{"role": "user", "content": "hi"}],
        },
        logging_obj=_build_logging_obj(CallTypes.aresponses.value, stream=True),
        model="gpt-5.4",
        args=(),
    )

    assert isinstance(result, CustomStreamWrapper)


def test_convert_cached_responses_bridge_chat_completion_nonstream():
    """openai/responses chat-completions bridge: non-streaming cache hit replays as ModelResponse."""
    from litellm import responses
    from litellm.types.utils import CallTypes, ModelResponse

    caching_handler = LLMCachingHandler(
        original_function=responses, request_kwargs={}, start_time=datetime.now()
    )
    cached_result = {
        "id": "chatcmpl-bridge-nonstream",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "gpt-5.4",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "Hi!"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18},
    }

    result = caching_handler._convert_cached_result_to_model_response(
        cached_result=cached_result,
        call_type=CallTypes.responses.value,
        kwargs={
            "model": "gpt-5.4",
            "stream": False,
            "messages": [{"role": "user", "content": "hi"}],
        },
        logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
        model="gpt-5.4",
        args=(),
    )

    assert isinstance(result, ModelResponse)
    assert result.choices[0].message.content == "Hi!"


def test_convert_cached_responses_legacy_nonstream_path():
    """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) falls through legacy path."""
    from litellm import responses
    from litellm.types.llms.openai import ResponsesAPIResponse
    from litellm.types.utils import CallTypes

    caching_handler = LLMCachingHandler(
        original_function=responses, request_kwargs={}, start_time=datetime.now()
    )
    cached_result = {
        "id": "resp_legacy_nonstream",
        "created_at": int(time.time()),
        "status": "completed",
        "model": "gpt-4o",
        "object": "response",
        "output": [
            {
                "type": "message",
                "id": "msg_legacy",
                "status": "completed",
                "role": "assistant",
                "content": [
                    {
                        "type": "output_text",
                        "text": "legacy response",
                        "annotations": [],
                    }
                ],
            }
        ],
    }

    result = caching_handler._convert_cached_result_to_model_response(
        cached_result=cached_result,
        call_type=CallTypes.responses.value,
        kwargs={"model": "gpt-4o", "input": "hi", "stream": False},
        logging_obj=_build_logging_obj(CallTypes.responses.value, stream=False),
        model="gpt-4o",
        args=(),
    )

    assert isinstance(result, ResponsesAPIResponse)
    assert result.id == "resp_legacy_nonstream"


def test_convert_cached_responses_legacy_stream_path():
    """Genuine ResponsesAPIResponse dict (no chatcmpl/choices) on stream falls through legacy path."""
    from litellm import responses
    from litellm.responses.streaming_iterator import (
        CachedResponsesAPIStreamingIterator,
    )
    from litellm.types.utils import CallTypes

    caching_handler = LLMCachingHandler(
        original_function=responses, request_kwargs={}, start_time=datetime.now()
    )
    cached_result = {
        "id": "resp_legacy_stream",
        "created_at": int(time.time()),
        "status": "completed",
        "model": "gpt-4o",
        "object": "response",
        "output": [
            {
                "type": "message",
                "id": "msg_legacy_stream",
                "status": "completed",
                "role": "assistant",
                "content": [
                    {
                        "type": "output_text",
                        "text": "legacy stream",
                        "annotations": [],
                    }
                ],
            }
        ],
    }

    result = caching_handler._convert_cached_result_to_model_response(
        cached_result=cached_result,
        call_type=CallTypes.responses.value,
        kwargs={"model": "gpt-4o", "input": "hi", "stream": True},
        logging_obj=_build_logging_obj(CallTypes.responses.value, stream=True),
        model="gpt-4o",
        args=(),
    )

    assert isinstance(result, CachedResponsesAPIStreamingIterator)


@pytest.mark.asyncio
async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input():
    """Image-embedding cache hit restores prompt_tokens=0 from the stored value
    instead of recomputing a bogus count by tokenizing the base64 input."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    # base64-like blob — token_counter over this would return a large nonzero count
    image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50

    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "amazon.titan-embed-image-v1",
            "prompt_tokens": 0,
            "prompt_tokens_details": {"image_count": 1},
        }
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="amazon.titan-embed-image-v1",
    )

    assert cache_hit
    assert response.usage is not None
    assert response.usage.prompt_tokens == 0
    assert response.usage.total_tokens == 0
    assert response.usage.prompt_tokens_details.image_count == 1


@pytest.mark.asyncio
async def test_embedding_cache_sums_stored_prompt_tokens_across_items():
    """A multi-item cache hit sums the stored per-item prompt_tokens back to the total."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    cached_result = [
        {
            "embedding": [-0.01],
            "index": 0,
            "object": "embedding",
            "model": "text-embedding-3-small",
            "prompt_tokens": 5,
        },
        {
            "embedding": [-0.02],
            "index": 1,
            "object": "embedding",
            "model": "text-embedding-3-small",
            "prompt_tokens": 4,
        },
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="text-embedding-3-small",
    )

    assert cache_hit
    assert response.usage.prompt_tokens == 9
    assert response.usage.total_tokens == 9


@pytest.mark.asyncio
async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries():
    """Legacy cache entries with no stored prompt_tokens still recompute via token_counter
    for str inputs (backward compatibility)."""
    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    # No prompt_tokens key — pre-fix entry
    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "text-embedding-ada-002",
        },
    ]

    mock_logging_obj = MagicMock()
    mock_logging_obj.async_success_handler = AsyncMock()
    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "text-embedding-ada-002", "input": "hello world"},
        logging_obj=mock_logging_obj,
        start_time=datetime.now(),
        model="text-embedding-ada-002",
    )

    assert cache_hit
    # token_counter over "hello world" yields a nonzero count — fallback path still runs
    assert response.usage.prompt_tokens > 0


@pytest.mark.asyncio
async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj():
    """A full embedding cache hit must stamp the resolved provider onto the logging
    obj so spend logs record the provider instead of None/unknown."""
    from litellm.types.utils import CallTypes

    llm_caching_handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs={},
        start_time=datetime.now(),
    )

    cached_result = [
        {
            "embedding": [-0.025, -0.019],
            "index": 0,
            "object": "embedding",
            "model": "text-embedding-3-small",
            "prompt_tokens": 5,
        }
    ]

    logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False)
    logging_obj.async_success_handler = AsyncMock()

    response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
        final_embedding_cached_response=None,
        cached_result=cached_result,
        kwargs={"model": "text-embedding-3-small", "input": "hello world"},
        logging_obj=logging_obj,
        start_time=datetime.now(),
        model="text-embedding-3-small",
    )

    assert cache_hit
    assert logging_obj.model_call_details["custom_llm_provider"] == "openai"


def test_request_kwargs_does_not_retain_logging_obj():
    """
    The caching handler lives on logging_obj._llm_caching_handler, so keeping
    litellm_logging_obj inside request_kwargs closes a reference cycle
    (Logging -> LLMCachingHandler -> kwargs -> Logging). That cycle keeps the
    full request payload alive until a generational GC pass instead of being
    freed by refcount when the request finishes; under bursts of large-token
    requests this presents as stepwise RSS growth that never returns to
    baseline. Other kwargs (messages included) must be preserved.
    """
    logging_obj = MagicMock()
    kwargs = {
        "model": "gpt-4o",
        "messages": [{"role": "user", "content": "hello"}],
        "litellm_logging_obj": logging_obj,
    }

    handler = LLMCachingHandler(
        original_function=MagicMock(),
        request_kwargs=kwargs,
        start_time=datetime.now(),
    )

    assert "litellm_logging_obj" not in handler.request_kwargs
    assert handler.request_kwargs["messages"] == kwargs["messages"]
    assert handler.request_kwargs["model"] == "gpt-4o"


def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
    """
    Regression test for the SDK losing async cache writes in short-lived scripts:
    async_set_cache dispatched the write as a bare fire-and-forget task, so
    asyncio.run cancelled it at loop close before the write landed (LIT-6184,
    deterministic with hiredis installed). The write must survive loop shutdown.
    """
    import litellm

    writes = []

    class _SlowWriteCache:
        supported_call_types = ["acompletion"]
        cache = None

        async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
            await asyncio.sleep(0.2)
            writes.append(result)

    async def acompletion(**kwargs):
        return None

    handler = LLMCachingHandler(
        original_function=acompletion,
        request_kwargs={},
        start_time=datetime.now(),
    )
    monkeypatch.setattr(litellm, "cache", _SlowWriteCache())

    async def _short_lived_script():
        await handler.async_set_cache(
            result=litellm.ModelResponse(),
            original_function=acompletion,
            kwargs={},
        )

    asyncio.run(_short_lived_script())

    assert len(writes) == 1


@pytest.mark.asyncio
async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monkeypatch):
    """The spend log for a cache hit must reuse the key the lookup already computed instead of hashing again."""
    import litellm
    from litellm.caching.caching import Cache
    from litellm.types.utils import CallTypes

    async def acompletion(**kwargs):
        return None

    monkeypatch.setattr(litellm, "cache", Cache(type="local"))
    kwargs = {"model": "gpt-5.4", "messages": [{"role": "user", "content": "hello"}], "caching": True}
    await litellm.cache.async_add_cache(
        litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}]), **kwargs
    )
    handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
    logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
    logging_obj.async_success_handler = AsyncMock()

    hit = await handler._async_get_cache(
        model="gpt-5.4",
        original_function=acompletion,
        logging_obj=logging_obj,
        start_time=datetime.now(),
        call_type=CallTypes.acompletion.value,
        kwargs=kwargs,
        args=(),
    )

    assert hit is not None and hit.cached_result is not None
    assert handler.preset_cache_key is not None
    assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
    assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key