mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
a173657dfb
* fix(caching): keep embedding cache hits aligned with request inputs Partial hits now send only the uncached inputs to the provider and merge fresh vectors back into their original positions. Responses whose item count differs from the input count (one input scoring many documents) are no longer written to the per-input cache, since a later hit would return a single item. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): drop mutable collection builds flagged by the type discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): bypass embedding cache entries written before the per input cardinality check Embedding cache entries now carry format_version and readers treat entries without it as misses, so entries that only hold the first row of a multi row response are refetched instead of served until their TTL expires. The provider call also receives a copy of the request kwargs with the uncached inputs rather than mutating the caller's mapping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(caching): assert a partial embedding cache hit becomes a full hit on repeat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(caching): await pending embedding cache writes before asserting on cache hits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): validate cached embeddings without mutating responses or request kwargs Validate cache rows through a frozen pydantic model so import does not depend on TypeAdapter support for ReadOnly TypedDicts, accept string embeddings, build the merged partial hit response instead of mutating the cached one, and hand the provider request mapping to post call hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): keep cache_hit and response_ms on merged partial embedding hits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
825 lines
29 KiB
Python
825 lines
29 KiB
Python
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 _PENDING_CACHE_WRITES, 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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
|
|
import litellm
|
|
from litellm.caching.caching import Cache
|
|
from litellm.types.utils import CallTypes
|
|
|
|
async def aanthropic_messages(**kwargs):
|
|
return None
|
|
|
|
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
|
kwargs = {
|
|
"model": "claude-sonnet-5",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"max_tokens": 16,
|
|
"caching": True,
|
|
"stream": False,
|
|
"_websearch_interception_converted_stream": True,
|
|
}
|
|
cached_message = {
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "hi"}],
|
|
}
|
|
await litellm.cache.async_add_cache(cached_message, **kwargs)
|
|
handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
|
|
logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
|
|
logging_obj.async_success_handler = AsyncMock()
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
|
|
|
hit = await handler._async_get_cache(
|
|
model="claude-sonnet-5",
|
|
original_function=aanthropic_messages,
|
|
logging_obj=logging_obj,
|
|
start_time=datetime.now(),
|
|
call_type=CallTypes.aanthropic_messages.value,
|
|
kwargs=kwargs,
|
|
args=(),
|
|
)
|
|
|
|
assert hit is not None and hit.cached_result == cached_message
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
|
|
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
|
|
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.6",
|
|
"messages": [{"role": "user", "content": "run the code"}],
|
|
"caching": True,
|
|
"stream": False,
|
|
"_code_interpreter_interception_converted_stream": True,
|
|
"_agentic_loop_depth": 1,
|
|
}
|
|
await litellm.cache.async_add_cache(
|
|
litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **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()
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
|
|
|
|
hit = await handler._async_get_cache(
|
|
model="gpt-5.6",
|
|
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 isinstance(hit.cached_result, litellm.ModelResponse)
|
|
assert hit.cached_result.choices[0].message.content == "done"
|
|
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
|
|
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_partial_embedding_cache_hit_sends_only_misses_and_keeps_input_order(monkeypatch):
|
|
import litellm
|
|
from litellm import CustomLLM
|
|
from litellm.caching.caching import Cache
|
|
from litellm.types.utils import Embedding, EmbeddingResponse
|
|
|
|
class RecordingEmbedder(CustomLLM):
|
|
provider_inputs: tuple[tuple[str, ...], ...] = ()
|
|
|
|
async def aembedding(self, model, input, model_response, **kwargs) -> EmbeddingResponse:
|
|
self.provider_inputs = (*self.provider_inputs, tuple(input))
|
|
return EmbeddingResponse(
|
|
model=model,
|
|
data=[
|
|
Embedding(embedding=[float(len(text))], index=idx, object="embedding")
|
|
for idx, text in enumerate(input)
|
|
],
|
|
)
|
|
|
|
embedder = RecordingEmbedder()
|
|
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "recording-embedder", "custom_handler": embedder}])
|
|
monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "recording-embedder"])
|
|
monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers, "recording-embedder"])
|
|
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
|
|
|
|
await litellm.aembedding(model="recording-embedder/m", input=["aa", "bbbb"])
|
|
await asyncio.gather(*_PENDING_CACHE_WRITES)
|
|
mixed_input = ["c", "aa", "ddd", "bbbb", "eeeee"]
|
|
response = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
|
|
await asyncio.gather(*_PENDING_CACHE_WRITES)
|
|
|
|
assert embedder.provider_inputs == (("aa", "bbbb"), ("c", "ddd", "eeeee")), embedder.provider_inputs
|
|
assert [item["index"] for item in response.data] == [0, 1, 2, 3, 4]
|
|
assert [item["embedding"] for item in response.data] == [[float(len(text))] for text in mixed_input]
|
|
assert response._hidden_params["cache_hit"] is True, "a partial hit must still be reported as a cache hit"
|
|
|
|
repeat = await litellm.aembedding(model="recording-embedder/m", input=mixed_input)
|
|
|
|
assert len(embedder.provider_inputs) == 2, embedder.provider_inputs
|
|
assert [item["embedding"] for item in repeat.data] == [[float(len(text))] for text in mixed_input]
|