mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
* fix(anthropic): support Bearer auth for custom api_base endpoints (Fixes #30926) * style: format common_utils.py with black * fix(anthropic): extract api_base from litellm_params in batches/files validate_environment * fix(anthropic): scope Bearer key check to custom api_base endpoints * fix(streaming): reset Anthropic message_start cursor (output_tokens=1) when no message_delta arrives The Anthropic streaming protocol emits `message_start.usage.output_tokens=1` as a placeholder cursor; the real cumulative output count only arrives in the final `message_delta` event. When a stream is cancelled before `message_delta` lands (common for thinking models on long-tail prompts), ChunkProcessor._calculate_usage_per_chunk's last-wins accumulator left completion_tokens stuck at 1. Because 1 is truthy, the `completion_tokens or token_counter(text=...)` fallback in calculate_usage() never fired, and requests were billed for 1 output token even when several thousand tokens of text had actually streamed. Fix: track whether any chunk's completion_tokens exceeded 1 (saw_non_cursor_completion). If the only update we saw was the cursor, reset completion_tokens to 0 so the text-based fallback estimates from the real completion content. Legitimate 1-token completions (model returns "Yes." etc.) are unaffected in practice — token_counter on a 1-token completion_output also yields ~1, so billing stays approximately correct. Tests: - TestAnthropicCursorBug (6 cases) — pins the post-fix behavior - TestNonAnthropicStreamingIntact (2 cases) — guards against regression on providers without the cursor pattern All 8 new tests pass; 9 existing streaming_chunk_builder_utils tests still pass. * fix(streaming): scope cursor reset to anthropic provider + recognize message_delta arrival Addresses both Greptile P2 threads on PR #30420: CLASS A — Anthropic-specific heuristic was applied globally ============================================================ The `completion_tokens == 1 and not saw_non_cursor_completion` reset lived in provider-neutral `streaming_chunk_builder_utils.py`. Any non-Anthropic provider that legitimately reports completion_tokens=1 in a single usage chunk (perfectly normal for short OpenAI / Bedrock / Vertex single-token replies with stream_options.include_usage=true) would have its value silently rewritten to 0 and re-billed via token_counter — producing a different number than what the provider actually charged. Fix: gate the reset on `custom_llm_provider == "anthropic"`, resolved from the first chunk's `_hidden_params` (the same field set by streaming_handler.py:722 on the live path). Unknown / missing provider is treated as non-Anthropic and skips the reset, so newer providers and custom plugins are also safe by default. CLASS B — `saw_non_cursor_completion` missed legitimate single-token replies ============================================================ Previous condition was `usage_chunk_dict["completion_tokens"] > 1`, which never fires for an Anthropic stream where the model legitimately emits exactly one output token (e.g., "Yes."). Anthropic still sends message_start (output_tokens=1, the cursor) AND message_delta (output_tokens=1, the real value) — same value, but two distinct usage events. The old check couldn't tell that apart from a cancelled stream where only message_start landed. Fix: track `completion_usage_updates` and flip `saw_non_cursor_completion` when EITHER (1) the value exceeds 1 (definitely not a placeholder), OR (2) we've seen >=2 completion-bearing usage events (positive evidence that message_delta arrived). Cancelled cursor-only streams still have exactly one event and still hit the reset; cache chunks with completion_tokens=0 don't count toward the threshold. Tests ============================================================ - _make_chunk now sets `_hidden_params["custom_llm_provider"]` (default "anthropic") so the gate is exercised by every existing test — none of them needed assertion changes besides the legitimate-single- token case, which now expects exactly 1 (was a fuzzy 0..3 range). - New: test_anthropic_cache_only_chunks_after_message_start_still_resets - New: test_non_anthropic_provider_completion_tokens_one_not_reset - New: test_unknown_provider_completion_tokens_one_not_reset 11/11 tests pass. * chore: add Co-authored-by trailer for attribution Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com> * fix(anthropic): preserve messages cache usage * style(anthropic): format messages cache usage helper * fix(anthropic): accept integral float cache token counts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(anthropic): accept integral float cache token counts * test(anthropic): cover cache usage edge cases * fix(gemini): preserve thoughtSignature for server-side tool responses When Gemini API returns toolCall and toolResponse parts, they might have different thoughtSignatures. Previously, LiteLLM merged them into a single dict, overwriting the response's thoughtSignature with the call's. This fix extracts them separately and re-injects them correctly. TAG=agy CONV=755b21d0-3200-40bc-bd1a-bb58a378a9a6 * fix(gemini): address PR comments on thoughtSignature handling - Fix orphan-response thoughtSignature regression by copying thought_signature to response_thought_signature - Add missing assertions in existing tests - Add new unit tests for orphan-response signature handling TAG=agy CONV=755b21d0-3200-40bc-bd1a-bb58a378a9a6 * feat(mcp): include server alias and server_id in mcp_info response - Add alias and server_id fields to mcp_info object in /mcp-rest/tools/list endpoint - Update rest_endpoints.py to surface alias from server config - Add test coverage in test_mcp_server.py and test_rest_endpoints.py Fixes #31015 * fix(proxy): reject non-finite spend via validate_finite_spend A NaN/-inf spend would bypass spend >= max_budget enforcement. Add a shared finite-value guard, defined above the litellm.proxy.* imports to avoid the module-level cyclic-import warning. * fix(proxy): require admin for any /key/update spend, reject non-finite Gate the admin check on the presence of `spend` (not a value diff): the DB spend lags the live cross-pod counter, so an "unchanged" spend on the non-admin path let a key owner / team member overwrite the live counter below real usage. Also reject NaN/+-inf spend before the DB write. * fix(proxy): invalidate spend counter on /user/update spend change A direct spend change on /user/update wrote the DB row but left the warm cross-pod counter at the stale value, so enforcement kept reading the old spend. Invalidate spend:user:{user_id} after the write (reseed-from-DB), and reject non-finite spend before the write. * fix(cache): route Bedrock semantic-cache sync embedding through the Router (#28244) The semantic cache's embedding model is a proxy Router alias whose AWS credentials (aws_role_name, aws_session_name) live only in the Router deployment's litellm_params. The sync embedding paths called litellm.embedding() directly, bypassing the Router, so they could neither resolve the alias nor assume the configured role; cross-account Bedrock semantic caching failed with "bedrock:InvokeModel is not authorized". On Redis this surfaced at proxy startup because redisvl's CustomTextVectorizer eagerly fires a dimension-probe embedding during cache construction, while llm_router is still None. Fix A: make the sync paths mirror the already-correct async paths. A shared, dependency-injected helper (litellm/caching/_embedding_router.py) decides whether to route through llm_router.embedding(...) when the model is a Router deployment, else fall back to direct litellm.embedding(...). Redis and qdrant sync set_cache/get_cache now precompute the embedding and pass vector= to the backend, exactly as the async astore/acheck already do. Both async _get_async_embedding methods are unified onto the same helper and now forward the caller's full metadata instead of a hand-picked subset. Fix B (Redis only): defer redisvl index construction from __init__ into a lazy, memoized llmcache property, so the dimension-probe embedding fires on first cache use, after llm_router is wired. A failed build is not memoized, so a transient outage recovers on the next request. Known limitation: resolve_embedding_router gates on an exact model-name match (same as the shipped async path); wildcard/alias/team-public routes still fall back to direct embedding. Tracked as a follow-up. * fix(cache): harden embedding-router and shrink Any surface (review) Address review feedback on the semantic-cache aws-role fix (#28244): - resolve_embedding_router now skips deployment entries missing model_name instead of raising KeyError on a malformed model_list (Greptile P2); add a regression test that fails on the old direct-key access. - Replace the `**kwargs: Any` passthrough on the four cache _get_embedding / _get_async_embedding helpers with an explicit, typed `metadata: Optional[Dict[str, Any]] = None` parameter. The helpers only ever consumed kwargs["metadata"], so this is behavior-preserving, makes the forwarded field obvious at the call site, and removes three bare-Any annotations (keeps the strict-rule ANN401 budget within ceiling). - Note in _build_llmcache that redisvl's dimension-probe embedding adds one extra billable embedding on the first cache request (Greptile P2). * fix(bedrock_mantle): correct responses routing for openai.gpt-5.x models Dashboard Test Connection for bedrock_mantle/openai.gpt-5.4 and openai.gpt-5.5 was failing with maximum recursion depth errors and "model does not exist" Route detection in the bedrock provider matched route tokens by plain substring, so the bedrock_mantle/ prefix was mistaken for the mantle/ invoke route and the body model was rewritten to bedrock_openai.gpt-5.5; route tokens now only match at a path-segment boundary so the bare model name is preserved A responses-mode model whose provider has no responses config bounced forever between the responses API and chat completions; the responses to completion fallback now tags its call so completion() does not bridge back, breaking the loop The Test Connection endpoint hardcoded the test mode to chat, which disabled mode auto-detection for responses-only models; the default is now None so the mode is detected from model capabilities acompletion() now drops a duplicate acompletion kwarg before building the partial and treats model_info=None as an empty dict to avoid a NoneType crash * test(bedrock_mantle): cover route guard and bridge flag; fix reportArgumentType regression Adds the regression coverage codecov flagged on the two responses to completion bridge guard lines and the bedrock route-prefix helper. The handler tests drive both the sync and async fallback paths with litellm.completion and litellm.acompletion mocked, and assert the forwarded kwargs carry _skip_responses_api_bridge=True, so dropping either flag line fails the suite. The common_utils tests assert that bedrock_mantle/openai.gpt-5.x no longer resolves to the mantle route while the genuine mantle/ and bedrock/mantle/ ids still do, exercising both branches of _model_has_route_prefix. Also aligns update_messages_with_model_file_ids model_id to Optional[str], matching its Responses API sibling, so the defensive model_info fallback no longer introduces a new reportArgumentType in completion(); the file-id lookup narrows model_id before the dict get * chore(ui): sync generated OpenAPI types for optional test_connection mode The test_model_connection mode body param default changed from chat to None so the mode is auto-detected from model capabilities, which makes the field optional in the proxy OpenAPI spec. Regenerate the committed schema so the dashboard types match: mode becomes optional and the description and default JSDoc follow the spec, keeping the Check UI API Types Sync gate green * refactor(bedrock): match all explicit route prefixes at path-segment boundary Migrates the remaining substring route checks to the existing _model_has_route_prefix helper so every explicit route token matches only as a leading path segment, consistent with get_bedrock_route and the mantle route. Covers _explicit_converse_route, _explicit_claude_platform_route, _explicit_invoke_route, _explicit_agent_route, _explicit_agentcore_route, _explicit_converse_like_route, _explicit_async_invoke_route and _explicit_openai_route. This also stops invoke/ from substring-matching async_invoke/. Route precedence and order are unchanged, and a note on the segment invariant is added to the helper docstring * test(bedrock): cover explicit route prefix segment matching Exercises all eight migrated _explicit_*_route helpers (converse, converse_like, invoke, async_invoke, agent, agentcore, claude_platform, openai) directly: each matches its token as a leading path segment and rejects the token glued to a preceding segment, so reverting any method to the old substring check fails the suite. Also asserts invoke/ no longer matches async_invoke/ models, the concrete improvement of the segment-boundary migration * test(proxy): assert negative spend is allowed (one-time grant use-case) Negative spend is intentionally permitted so admins can grant extra allowance for the current budget period only, without raising the recurring budget ceiling. Cover it explicitly in validate_finite_spend and via the /user/update invalidation test. * fix(google_genai): forward native generateContent top-level fields Google's native generateContent REST body carries safetySettings, toolConfig, cachedContent and labels at the top level as siblings of generationConfig. The proxy's :generateContent endpoint spread them into agenerate_content as loose kwargs and then dropped them, so callers had to wrap them in extra_body for them to take effect; safetySettings, for instance, was silently ignored The provider config now exposes the native top-level field names and setup_generate_content_call collects whichever are present, merging them into the outgoing request body through the existing extra_body merge so they reach Google verbatim. An explicit extra_body still wins on conflict. The sync generate_content_stream path now also forwards systemInstruction, matching the other three entry points Fixes #12671 Claude-Session: https://claude.ai/code/session_016MFtMXokCjT8u6mvyASudK * fix(proxy): resolve env refs for DB-stored models * fix(proxy): restrict DB env ref resolution * fix(proxy): block team DB env ref resolution * fix(lint): resolve ANN401/UP045/C901 strict-gate violations - Replace Optional[X] with X | None (UP045) in 8 files - Replace Any return/param types with concrete types or object (ANN401) - Extract _make_api_key_auth_header helper to reduce get_anthropic_headers complexity below C901 threshold (17 → 14) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(anthropic): preserve x-api-key for custom endpoints; opt-in Bearer via prefix Users who pass a key already prefixed with "Bearer " get Authorization: Bearer. All other keys continue to use x-api-key, preserving backward compatibility with custom api_base endpoints that expect x-api-key rather than Authorization. Also consolidates get_auth_header to reuse _make_api_key_auth_header helper, eliminating the duplicated custom-endpoint routing logic. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * revert(anthropic): restore Bearer routing for non-sk-ant- keys on custom api_base The backwards-compat change broke existing tests that verify the intentional Bearer-for-custom-base behavior (Fixes #30926). Restore original logic while keeping the _make_api_key_auth_header helper for code deduplication. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(anthropic): gate Bearer-for-custom-base behind use_bearer_for_custom_base flag Previously the auth-header switch from x-api-key to Authorization: Bearer applied unconditionally for non-sk-ant- keys on a custom api_base, silently breaking existing deployments that proxied to gateways expecting x-api-key. Introduce use_bearer_for_custom_base: bool = False on _make_api_key_auth_header, get_anthropic_headers, and get_auth_header. validate_environment reads it from litellm_params so callers can opt in per-model without any API surface change. Tests updated to pass use_bearer_for_custom_base=True where Bearer behavior is asserted. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(redis): apply namespace prefix in delete_cache and async_delete_cache (#29981) DEL was the only Redis cache operation that skipped check_and_fix_namespace, so it targeted the raw SHA256 hash (e.g. 3997c4...) rather than the namespaced key (litellm:3997c4...). This caused two problems: a Redis NOPERM error on deployments with an ACL restricting DEL to the litellm:* pattern, and a silent no-op on all other deployments since the un-prefixed key was never stored. * style(anthropic): reformat common_utils.py with Black (--target-version py312) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: preserve cache metadata and spend counters * style: apply ruff format to streaming_iterator.py * refactor: reduce complexity of usage/spend helpers to satisfy strict ruff gate Extract Anthropic message_start cursor reset into _reset_anthropic_cursor_completion_tokens and the cross-pod spend-counter invalidation into _invalidate_user_spend_counter_if_changed, keeping both _calculate_usage_per_chunk and _update_single_user_helper under the max-complexity ceiling. Use builtin generics in the new signatures so no new UP006 violations are introduced. Behavior unchanged. --------- Co-authored-by: rupak-eng <rupakji99@gmail.com> Co-authored-by: songkuan-zheng <252822057+songkuan-zheng@users.noreply.github.com> Co-authored-by: songkuan-zheng <songkuan-zheng@users.noreply.github.com> Co-authored-by: Kannan Priyadharshan <kpd2204@gmail.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Marco Georgaklis <mgeorgaklis@google.com> Co-authored-by: Anjaiah Methuku <anjaiahspr@gmail.com> Co-authored-by: Andrii Butko <booandrew23@gmail.com> Co-authored-by: Kent <kingdooo@gmail.com> Co-authored-by: kunal2002 <k.nayyar2002@gmail.com> Co-authored-by: Ali Khan <alirazakhan.offi@gmail.com> Co-authored-by: jesco-absolut <team@srswti.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Matt Hill <mhill@dataminr.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>
1235 lines
42 KiB
Python
1235 lines
42 KiB
Python
import os
|
|
import sys
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
|
|
# Tests for RedisSemanticCache
|
|
def test_redis_semantic_cache_initialization(monkeypatch):
|
|
# Mock the redisvl import
|
|
semantic_cache_mock = MagicMock()
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(CustomTextVectorizer=MagicMock()),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize the cache with a similarity threshold
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Verify the semantic cache was initialized with correct parameters
|
|
assert redis_semantic_cache.similarity_threshold == 0.8
|
|
|
|
# Use pytest.approx for floating point comparison to handle precision issues
|
|
assert redis_semantic_cache.distance_threshold == pytest.approx(0.2, abs=1e-10)
|
|
assert redis_semantic_cache.embedding_model == "text-embedding-ada-002"
|
|
|
|
# Test initialization with missing similarity_threshold
|
|
with pytest.raises(ValueError, match="similarity_threshold must be provided"):
|
|
RedisSemanticCache()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache(monkeypatch):
|
|
# Mock the redisvl import and embedding function
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize cache
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Mock the llmcache.check method to return a result
|
|
mock_result = [
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris is the capital of France."}',
|
|
"vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
redis_semantic_cache.llmcache.check = MagicMock(return_value=mock_result)
|
|
|
|
# Mock the embedding function
|
|
with (
|
|
patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
),
|
|
patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
),
|
|
):
|
|
# Test get_cache with a message
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata=metadata,
|
|
)
|
|
|
|
# Verify result is properly parsed
|
|
assert result == {"content": "Paris is the capital of France."}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
|
|
# Verify llmcache.check was called
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
),
|
|
patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
),
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.store = MagicMock()
|
|
|
|
with patch(
|
|
"litellm.embedding",
|
|
return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]},
|
|
):
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
ttl=60,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
ttl=60,
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch):
|
|
fallback_cache_mock = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("stored index schema differs from requested fields"),
|
|
fallback_cache_mock,
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
|
|
assert redis_semantic_cache.llmcache is fallback_cache_mock
|
|
assert semantic_cache_mock.call_args_list[0].kwargs["name"] == "existing_index"
|
|
assert (
|
|
semantic_cache_mock.call_args_list[1].kwargs["name"]
|
|
== "existing_index_isolated"
|
|
)
|
|
assert semantic_cache_mock.call_args_list[1].kwargs["filterable_fields"] == [
|
|
RedisSemanticCache._cache_key_filterable_field()
|
|
]
|
|
|
|
|
|
def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch):
|
|
fallback_cache_mock = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("Existing index schema does not match"),
|
|
ValueError("Existing index schema does not match"),
|
|
fallback_cache_mock,
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
|
|
assert redis_semantic_cache.llmcache is fallback_cache_mock
|
|
assert (
|
|
semantic_cache_mock.call_args_list[2].kwargs["name"]
|
|
== "existing_index_isolated"
|
|
)
|
|
assert semantic_cache_mock.call_args_list[2].kwargs["overwrite"] is True
|
|
assert semantic_cache_mock.call_args_list[2].kwargs["filterable_fields"] == [
|
|
RedisSemanticCache._cache_key_filterable_field()
|
|
]
|
|
|
|
|
|
def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypatch):
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[
|
|
ValueError("Existing index schema does not match"),
|
|
ValueError("connection failed"),
|
|
]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
with pytest.raises(ValueError, match="connection failed"):
|
|
cache = RedisSemanticCache(
|
|
similarity_threshold=0.8,
|
|
index_name="existing_index",
|
|
)
|
|
_ = cache.llmcache
|
|
|
|
|
|
def test_redis_semantic_cache_reraises_unexpected_index_error():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.distance_threshold = 0.2
|
|
semantic_cache_mock = MagicMock(side_effect=ValueError("connection failed"))
|
|
|
|
with pytest.raises(ValueError, match="connection failed"):
|
|
redis_semantic_cache._init_semantic_cache(
|
|
semantic_cache_cls=semantic_cache_mock,
|
|
index_name="existing_index",
|
|
redis_url="redis://localhost:6379",
|
|
cache_vectorizer=MagicMock(),
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_matches_bytes_cache_key():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
assert redis_semantic_cache._cache_hit_matches_key(
|
|
cache_hit={RedisSemanticCache.CACHE_KEY_FIELD_NAME: b"test_key"},
|
|
key="test_key",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_rejects_pre_isolation_unscoped_hit():
|
|
"""Pre-isolation entries with no cache-key field cannot be safely
|
|
reassigned to a caller's scope and are treated as misses."""
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
cache_hit = {
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
assert not redis_semantic_cache._cache_hit_matches_key(
|
|
cache_hit=cache_hit,
|
|
key="test_key",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_builds_filter_expression(monkeypatch):
|
|
class FakeTag:
|
|
def __init__(self, field_name):
|
|
self.field_name = field_name
|
|
|
|
def __eq__(self, value):
|
|
return (self.field_name, value)
|
|
|
|
with patch.dict("sys.modules", {"redisvl.query.filter": MagicMock(Tag=FakeTag)}):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
|
|
assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == (
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME,
|
|
"test_key",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_get_cache(monkeypatch):
|
|
# Mock the redisvl import
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
# Set environment variables
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
# Initialize cache
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
# Mock the async methods
|
|
mock_result = [
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris is the capital of France."}',
|
|
"vector_distance": 0.1, # Distance of 0.1 means similarity of 0.9
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=mock_result)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
# Test async_get_cache with a message
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
# Verify result is properly parsed
|
|
assert result == {"content": "Paris is the capital of France."}
|
|
|
|
# Verify methods were called
|
|
redis_semantic_cache._get_async_embedding.assert_called_once()
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
assert result is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter(
|
|
monkeypatch,
|
|
):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
redis_semantic_cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(
|
|
return_value=[0.1, 0.2, 0.3]
|
|
)
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
ttl=60,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
ttl=60,
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_uses_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_uses_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.similarity_threshold = 0.8
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_flattens_structured_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=[
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "input_text", "text": "What is the capital of France?"},
|
|
{"type": "input_text", "text": "Answer briefly."},
|
|
{
|
|
"type": "input_image",
|
|
"image_url": "https://example.com/paris.png",
|
|
},
|
|
],
|
|
}
|
|
],
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?\nAnswer briefly.",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_prefers_messages():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
messages=[{"content": "message prompt"}],
|
|
input="responses prompt",
|
|
)
|
|
|
|
assert prompt == "message prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_handles_model_objects():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ModelDumpInput:
|
|
def model_dump(self):
|
|
return {"content": [{"text": "model dump prompt"}]}
|
|
|
|
class DictInput:
|
|
def dict(self):
|
|
return {"content": [{"output_text": "dict prompt"}]}
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
input=[
|
|
ModelDumpInput(),
|
|
DictInput(),
|
|
{"content": [{"input_text": "inline prompt"}]},
|
|
{"content": [{"type": "input_image", "image_url": "https://example.com"}]},
|
|
]
|
|
)
|
|
|
|
assert prompt == "model dump prompt\ndict prompt\ninline prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_returns_none_without_text():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
assert RedisSemanticCache._get_prompt_from_kwargs() is None
|
|
assert RedisSemanticCache._get_prompt_from_kwargs(input=None) is None
|
|
assert RedisSemanticCache._get_prompt_from_kwargs(input=" ") is None
|
|
assert (
|
|
RedisSemanticCache._get_prompt_from_kwargs(
|
|
input=[{"type": "input_image", "image_url": "https://example.com"}]
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_skips_blank_dict_text_keys():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(
|
|
input={"text": " ", "input_text": "fallback prompt"}
|
|
)
|
|
|
|
assert prompt == "fallback prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_skips_blank_object_text_keys():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ResponseInput:
|
|
text = " "
|
|
input_text = "fallback prompt"
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput())
|
|
|
|
assert prompt == "fallback prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_prompt_extraction_handles_object_content():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
class ResponseInput:
|
|
content = [{"text": "object content prompt"}]
|
|
|
|
prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput())
|
|
|
|
assert prompt == "object content prompt"
|
|
|
|
|
|
def test_redis_semantic_cache_set_cache_skips_blank_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
|
|
redis_semantic_cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=" ",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.store.assert_not_called()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_sets_similarity_on_blank_responses_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
metadata = {}
|
|
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input=" ",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.check.assert_not_called()
|
|
|
|
|
|
def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.check = MagicMock(return_value=[])
|
|
redis_semantic_cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = redis_semantic_cache.get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
assert result is None
|
|
assert metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_paths_use_responses_string_input():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.similarity_threshold = 0.8
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
redis_semantic_cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
redis_semantic_cache._get_ttl = MagicMock(return_value=None)
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
metadata = {}
|
|
result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
assert result == {"content": "Paris"}
|
|
assert metadata["semantic-similarity"] == pytest.approx(0.9)
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_semantic_cache_async_paths_set_similarity_on_misses():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
redis_semantic_cache.llmcache = MagicMock()
|
|
redis_semantic_cache.llmcache.astore = AsyncMock()
|
|
redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=[])
|
|
redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
await redis_semantic_cache.async_set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
input=" ",
|
|
)
|
|
|
|
redis_semantic_cache.llmcache.astore.assert_not_called()
|
|
redis_semantic_cache._get_async_embedding.assert_not_called()
|
|
|
|
blank_metadata = {}
|
|
blank_result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input=" ",
|
|
metadata=blank_metadata,
|
|
)
|
|
|
|
assert blank_result is None
|
|
assert blank_metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.acheck.assert_not_called()
|
|
redis_semantic_cache._get_async_embedding.assert_not_called()
|
|
|
|
with patch.object(
|
|
redis_semantic_cache,
|
|
"_get_cache_key_filter_expression",
|
|
return_value="cache-key-filter",
|
|
):
|
|
miss_metadata = {}
|
|
miss_result = await redis_semantic_cache.async_get_cache(
|
|
key="test_key",
|
|
input="What is the capital of France?",
|
|
metadata=miss_metadata,
|
|
)
|
|
|
|
assert miss_result is None
|
|
assert miss_metadata["semantic-similarity"] == 0.0
|
|
redis_semantic_cache.llmcache.acheck.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
def test_redis_get_embedding_routes_through_router(monkeypatch):
|
|
import sys
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
|
|
router = MagicMock()
|
|
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
with patch("litellm.embedding") as direct_embed:
|
|
vec = cache._get_embedding("hello", metadata={"user_api_key": "sk-x"})
|
|
|
|
assert vec == [0.5, 0.6]
|
|
router.embedding.assert_called_once()
|
|
assert router.embedding.call_args.kwargs["model"] == "sem-embed"
|
|
assert router.embedding.call_args.kwargs["input"] == "hello"
|
|
assert router.embedding.call_args.kwargs["cache"] == {
|
|
"no-store": True,
|
|
"no-cache": True,
|
|
}
|
|
assert router.embedding.call_args.kwargs["metadata"] == {
|
|
"user_api_key": "sk-x",
|
|
"semantic-cache-embedding": True,
|
|
}
|
|
direct_embed.assert_not_called()
|
|
|
|
|
|
def test_redis_get_embedding_falls_back_to_direct(monkeypatch):
|
|
import sys
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "text-embedding-ada-002"
|
|
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = None
|
|
fake_proxy.llm_model_list = None
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
with patch(
|
|
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2]}]}
|
|
) as direct_embed:
|
|
vec = cache._get_embedding("hello")
|
|
|
|
assert vec == [0.1, 0.2]
|
|
direct_embed.assert_called_once()
|
|
|
|
|
|
def test_cache_get_cache_passes_responses_input_to_backend_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.cache.get_cache = MagicMock(return_value=None)
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
|
|
metadata = {}
|
|
cache.get_cache(
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={},
|
|
)
|
|
|
|
cache.cache.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_filters_non_lookup_kwargs_from_backend_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache._get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
|
|
def _cache_hit(_cache_key, **cache_kwargs):
|
|
cache_kwargs["metadata"]["semantic-similarity"] = 0.7
|
|
return {"content": "Paris"}
|
|
|
|
cache.cache.get_cache = MagicMock(side_effect=_cache_hit)
|
|
|
|
metadata = {"user_api_key": "sk-secret", "trace_id": "trace-id"}
|
|
result = cache.get_cache(
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={"s-maxage": 10},
|
|
api_key="sk-secret",
|
|
headers={"authorization": "Bearer sk-secret"},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
assert metadata == {
|
|
"user_api_key": "sk-secret",
|
|
"trace_id": "trace-id",
|
|
"semantic-similarity": 0.7,
|
|
}
|
|
|
|
forwarded_kwargs = cache.cache.get_cache.call_args.kwargs
|
|
assert forwarded_kwargs == {
|
|
"input": "What is the capital of France?",
|
|
"metadata": {
|
|
"user_api_key": "sk-secret",
|
|
"trace_id": "trace-id",
|
|
"semantic-similarity": 0.7,
|
|
},
|
|
}
|
|
assert forwarded_kwargs["metadata"] is not metadata
|
|
cache._get_cache_logic.assert_called_once_with(
|
|
cached_result={"content": "Paris"},
|
|
max_age=10,
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_filters_sensitive_kwargs_without_metadata():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.cache = MagicMock()
|
|
cache.cache.get_cache = MagicMock(return_value={"content": "Paris"})
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache._get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
|
|
result = cache.get_cache(
|
|
input="What is the capital of France?",
|
|
cache={"s-maxage": 10},
|
|
api_key="sk-secret",
|
|
headers={"authorization": "Bearer sk-secret"},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
cache.cache.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
)
|
|
|
|
|
|
def test_cache_get_cache_passes_responses_input_to_dynamic_cache():
|
|
from litellm.caching.caching import Cache
|
|
|
|
cache = Cache.__new__(Cache)
|
|
cache.should_use_cache = MagicMock(return_value=True)
|
|
cache.get_cache_key = MagicMock(return_value="test_key")
|
|
cache._get_cache_logic = MagicMock(return_value={"content": "Paris"})
|
|
dynamic_cache_object = MagicMock()
|
|
dynamic_cache_object.get_cache = MagicMock(return_value={"content": "Paris"})
|
|
|
|
metadata = {}
|
|
result = cache.get_cache(
|
|
dynamic_cache_object=dynamic_cache_object,
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
cache={},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
dynamic_cache_object.get_cache.assert_called_once_with(
|
|
"test_key",
|
|
input="What is the capital of France?",
|
|
metadata=metadata,
|
|
)
|
|
cache._get_cache_logic.assert_called_once_with(
|
|
cached_result={"content": "Paris"},
|
|
max_age=float("inf"),
|
|
)
|
|
|
|
|
|
def test_redis_sync_set_cache_passes_precomputed_vector():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.llmcache = MagicMock()
|
|
cache._get_cache_filters = MagicMock(
|
|
return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}
|
|
)
|
|
cache._get_ttl = MagicMock(return_value=None)
|
|
cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
cache.set_cache(
|
|
key="test_key",
|
|
value={"content": "Paris"},
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
)
|
|
|
|
cache._get_embedding.assert_called_once()
|
|
cache.llmcache.store.assert_called_once_with(
|
|
"What is the capital of France?",
|
|
"{'content': 'Paris'}",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"},
|
|
)
|
|
|
|
|
|
def test_redis_sync_get_cache_passes_precomputed_vector():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.similarity_threshold = 0.8
|
|
cache.llmcache = MagicMock()
|
|
cache.llmcache.check = MagicMock(
|
|
return_value=[
|
|
{
|
|
"prompt": "What is the capital of France?",
|
|
"response": '{"content": "Paris"}',
|
|
"vector_distance": 0.1,
|
|
RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
|
|
}
|
|
]
|
|
)
|
|
cache._get_embedding = MagicMock(return_value=[0.1, 0.2, 0.3])
|
|
|
|
with patch.object(
|
|
cache, "_get_cache_key_filter_expression", return_value="cache-key-filter"
|
|
):
|
|
result = cache.get_cache(
|
|
key="test_key",
|
|
messages=[{"content": "What is the capital of France?"}],
|
|
metadata={},
|
|
)
|
|
|
|
assert result == {"content": "Paris"}
|
|
cache._get_embedding.assert_called_once()
|
|
cache.llmcache.check.assert_called_once_with(
|
|
prompt="What is the capital of France?",
|
|
vector=[0.1, 0.2, 0.3],
|
|
filter_expression="cache-key-filter",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_async_embedding_forwards_full_metadata(monkeypatch):
|
|
import sys
|
|
import types
|
|
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
cache.embedding_model = "sem-embed"
|
|
|
|
router = MagicMock()
|
|
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
|
|
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
|
|
fake_proxy.llm_router = router
|
|
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
|
|
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
|
|
|
|
await cache._get_async_embedding(
|
|
"hello",
|
|
metadata={"user_api_key": "sk-x", "user_api_key_team_id": "team-1"},
|
|
)
|
|
|
|
md = router.aembedding.call_args.kwargs["metadata"]
|
|
assert md["user_api_key"] == "sk-x"
|
|
assert md["user_api_key_team_id"] == "team-1" # FAILS today: team_id is dropped
|
|
assert md["semantic-cache-embedding"] is True
|
|
|
|
|
|
def test_redis_init_defers_redisvl_construction(monkeypatch):
|
|
semantic_cache_mock = MagicMock()
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
semantic_cache_mock.assert_not_called()
|
|
custom_vectorizer_mock.assert_not_called()
|
|
|
|
first = cache.llmcache
|
|
semantic_cache_mock.assert_called_once()
|
|
custom_vectorizer_mock.assert_called_once()
|
|
|
|
second = cache.llmcache
|
|
assert first is second
|
|
semantic_cache_mock.assert_called_once()
|
|
|
|
|
|
def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch):
|
|
built_cache = MagicMock()
|
|
semantic_cache_mock = MagicMock(
|
|
side_effect=[ConnectionError("redis down"), built_cache]
|
|
)
|
|
custom_vectorizer_mock = MagicMock()
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock),
|
|
"redisvl.utils.vectorize": MagicMock(
|
|
CustomTextVectorizer=custom_vectorizer_mock
|
|
),
|
|
},
|
|
):
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
monkeypatch.setenv("REDIS_HOST", "localhost")
|
|
monkeypatch.setenv("REDIS_PORT", "6379")
|
|
monkeypatch.setenv("REDIS_PASSWORD", "test_password")
|
|
|
|
cache = RedisSemanticCache(similarity_threshold=0.8)
|
|
|
|
with pytest.raises(ConnectionError, match="redis down"):
|
|
_ = cache.llmcache
|
|
|
|
assert cache.llmcache is built_cache
|
|
assert semantic_cache_mock.call_count == 2
|
|
|
|
|
|
def test_redis_llmcache_setter_supported():
|
|
from litellm.caching.redis_semantic_cache import RedisSemanticCache
|
|
|
|
cache = RedisSemanticCache.__new__(RedisSemanticCache)
|
|
sentinel = MagicMock()
|
|
cache.llmcache = sentinel
|
|
assert cache.llmcache is sentinel
|