import asyncio import json import os import sys import time from unittest.mock import MagicMock, patch import httpx import pytest import respx from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path from datetime import datetime from unittest.mock import AsyncMock, MagicMock 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)