litellm/tests/test_litellm/caching/test_qdrant_semantic_cache.py
Sameer Kankute 7eacdd5258
chore: litellm oss staging 250626 (#31305)
* 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>
2026-06-25 21:00:28 -07:00

910 lines
33 KiB
Python

import os
import sys
import types
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
def test_qdrant_semantic_cache_initialization(monkeypatch):
"""
Test QDRANT semantic cache initialization with proper parameters.
Verifies that the cache is initialized correctly with given configuration.
"""
# Mock the httpx clients and API calls
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_index_response = MagicMock()
mock_index_response.status_code = 200
mock_sync_client_instance.put.return_value = mock_index_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize the cache with similarity threshold
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Verify the cache was initialized with correct parameters
assert qdrant_cache.collection_name == "test_collection"
assert qdrant_cache.qdrant_api_base == "http://test.qdrant.local"
assert qdrant_cache.qdrant_api_key == "test_key"
assert qdrant_cache.similarity_threshold == 0.8
mock_sync_client_instance.put.assert_called_once_with(
url="http://test.qdrant.local/collections/test_collection/index",
headers={
"Content-Type": "application/json",
"api-key": "test_key",
},
json={
"field_name": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"field_schema": "keyword",
},
)
# Test initialization with missing similarity_threshold
with pytest.raises(Exception, match="similarity_threshold must be provided"):
QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
)
def test_qdrant_semantic_cache_get_cache_hit():
"""
Test QDRANT semantic cache get method when there's a cache hit.
Verifies that cached results are properly retrieved and parsed.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock a cache hit result from search API
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {
"result": [
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of France?", # Original prompt
"response": '{"id": "test-123", "choices": [{"message": {"content": "Paris is the capital of France."}}]}',
},
"score": 0.9,
}
]
}
qdrant_cache.sync_client.post = MagicMock(return_value=mock_search_response)
# Mock the embedding function
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
# Test get_cache with a message
result = qdrant_cache.get_cache(
key="test_key", messages=[{"content": "What is the capital of France?"}]
)
# Verify result is properly parsed
expected_result = {
"id": "test-123",
"choices": [
{"message": {"content": "Paris is the capital of France."}}
],
}
assert result == expected_result
# Verify search was called
qdrant_cache.sync_client.post.assert_called()
assert qdrant_cache.sync_client.post.call_args.kwargs["json"]["filter"] == {
"must": [
{
"key": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"match": {"value": "test_key"},
}
]
}
def test_qdrant_semantic_cache_rejects_unscoped_cache_hit():
"""
Test QDRANT semantic cache rejects old or unscoped cache hits.
Legacy points have only text and response payloads, so they cannot be
safely migrated to a generated LiteLLM cache key.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {
"result": [
{
"payload": {
"text": "What is the capital of France?",
"response": '{"id": "test-123"}',
},
"score": 0.9,
}
]
}
qdrant_cache.sync_client.post = MagicMock(return_value=mock_search_response)
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
metadata = {}
result = qdrant_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_qdrant_semantic_cache_payload_index_failure_is_non_blocking():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {"Content-Type": "application/json"}
qdrant_cache.sync_client = MagicMock()
response = MagicMock()
response.status_code = 400
response.text = "bad index"
qdrant_cache.sync_client.put.return_value = response
qdrant_cache._ensure_cache_key_payload_index()
qdrant_cache.sync_client.put.assert_called_once()
def test_qdrant_semantic_cache_payload_index_exception_is_non_blocking():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {"Content-Type": "application/json"}
qdrant_cache.sync_client = MagicMock()
qdrant_cache.sync_client.put.side_effect = Exception("boom")
qdrant_cache._ensure_cache_key_payload_index()
qdrant_cache.sync_client.put.assert_called_once()
def _mock_qdrant_get_cache_result(qdrant_result):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
qdrant_cache.embedding_model = "text-embedding-ada-002"
qdrant_cache.qdrant_api_base = "http://test.qdrant.local"
qdrant_cache.collection_name = "test_collection"
qdrant_cache.headers = {
"Content-Type": "application/json",
"api-key": "test_key",
}
qdrant_cache.similarity_threshold = 0.8
qdrant_cache.sync_client = MagicMock()
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {"result": qdrant_result}
qdrant_cache.sync_client.post.return_value = mock_search_response
return qdrant_cache, QdrantSemanticCache
@pytest.mark.parametrize("qdrant_result", [None, []])
def test_qdrant_semantic_cache_get_cache_sets_metadata_on_empty_miss(qdrant_result):
qdrant_cache, _ = _mock_qdrant_get_cache_result(qdrant_result)
metadata = {}
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
result = qdrant_cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of Spain?"}],
metadata=metadata,
)
assert result is None
assert metadata["semantic-similarity"] == 0.0
def test_qdrant_semantic_cache_get_cache_sets_metadata_on_below_threshold_miss():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
qdrant_cache, _ = _mock_qdrant_get_cache_result(
[
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of Spain?",
"response": '{"id": "test-456"}',
},
"score": 0.7,
}
]
)
metadata = {}
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
result = qdrant_cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of Spain?"}],
metadata=metadata,
)
assert result is None
assert metadata["semantic-similarity"] == 0.7
def test_qdrant_semantic_cache_get_cache_miss():
"""
Test QDRANT semantic cache get method when there's a cache miss.
Verifies that None is returned when no similar cached results are found.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock a cache miss (no results)
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {"result": []}
qdrant_cache.sync_client.post = MagicMock(return_value=mock_search_response)
# Mock the embedding function
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.2, 0.3]}]}
):
# Test get_cache with a message
result = qdrant_cache.get_cache(
key="test_key", messages=[{"content": "What is the capital of Spain?"}]
)
# Verify None is returned for cache miss
assert result is None
# Verify search was called
qdrant_cache.sync_client.post.assert_called()
@pytest.mark.asyncio
async def test_qdrant_semantic_cache_async_get_cache_hit():
"""
Test QDRANT semantic cache async get method when there's a cache hit.
Verifies that cached results are properly retrieved and parsed asynchronously.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
# Mock async client
mock_async_client_instance = AsyncMock()
mock_async_client.return_value = mock_async_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock a cache hit result from async search API
# Note: .json() should be sync even for async responses
mock_search_response = MagicMock()
mock_search_response.status_code = 200
mock_search_response.json.return_value = {
"result": [
{
"payload": {
QdrantSemanticCache.CACHE_KEY_FIELD_NAME: "test_key",
"text": "What is the capital of Spain?", # Original prompt
"response": '{"id": "test-456", "choices": [{"message": {"content": "Madrid is the capital of Spain."}}]}',
},
"score": 0.85,
}
]
}
qdrant_cache.async_client.post = AsyncMock(return_value=mock_search_response)
# Mock the async embedding function
with patch(
"litellm.aembedding",
return_value={"data": [{"embedding": [0.4, 0.5, 0.6]}]},
):
# Test async_get_cache with a message
result = await qdrant_cache.async_get_cache(
key="test_key",
messages=[{"content": "What is the capital of Spain?"}],
metadata={},
)
# Verify result is properly parsed
expected_result = {
"id": "test-456",
"choices": [
{"message": {"content": "Madrid is the capital of Spain."}}
],
}
assert result == expected_result
# Verify async search was called
qdrant_cache.async_client.post.assert_called()
assert qdrant_cache.async_client.post.call_args.kwargs["json"][
"filter"
] == {
"must": [
{
"key": QdrantSemanticCache.CACHE_KEY_FIELD_NAME,
"match": {"value": "test_key"},
}
]
}
@pytest.mark.asyncio
async def test_qdrant_semantic_cache_async_get_cache_miss():
"""
Test QDRANT semantic cache async get method when there's a cache miss.
Verifies that None is returned when no similar cached results are found.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
# Mock async client
mock_async_client_instance = AsyncMock()
mock_async_client.return_value = mock_async_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock a cache miss (no results)
mock_search_response = MagicMock() # Note: .json() should be sync
mock_search_response.status_code = 200
mock_search_response.json.return_value = {"result": []}
qdrant_cache.async_client.post = AsyncMock(return_value=mock_search_response)
# Mock the async embedding function
with patch(
"litellm.aembedding",
return_value={"data": [{"embedding": [0.7, 0.8, 0.9]}]},
):
# Test async_get_cache with a message
result = await qdrant_cache.async_get_cache(
key="test_key",
messages=[{"content": "What is the capital of Italy?"}],
metadata={},
)
# Verify None is returned for cache miss
assert result is None
# Verify async search was called
qdrant_cache.async_client.post.assert_called()
def test_qdrant_semantic_cache_set_cache():
"""
Test QDRANT semantic cache set method.
Verifies that responses are properly stored in the cache.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock the upsert method
mock_upsert_response = MagicMock()
mock_upsert_response.status_code = 200
qdrant_cache.sync_client.put = MagicMock(return_value=mock_upsert_response)
# Mock response to cache
response_to_cache = {
"id": "test-789",
"choices": [{"message": {"content": "Rome is the capital of Italy."}}],
}
# Mock the embedding function
with patch(
"litellm.embedding", return_value={"data": [{"embedding": [0.1, 0.1, 0.1]}]}
):
# Test set_cache
qdrant_cache.set_cache(
key="test_key",
value=response_to_cache,
messages=[{"content": "What is the capital of Italy?"}],
)
# Verify upsert was called
qdrant_cache.sync_client.put.assert_called()
upsert_payload = qdrant_cache.sync_client.put.call_args.kwargs["json"][
"points"
][0]["payload"]
assert (
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
)
@pytest.mark.asyncio
async def test_qdrant_semantic_cache_async_set_cache():
"""
Test QDRANT semantic cache async set method.
Verifies that responses are properly stored in the cache asynchronously.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
# Mock async client
mock_async_client_instance = AsyncMock()
mock_async_client.return_value = mock_async_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize cache
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Mock the async upsert method
mock_upsert_response = MagicMock() # Note: .json() should be sync
mock_upsert_response.status_code = 200
qdrant_cache.async_client.put = AsyncMock(return_value=mock_upsert_response)
# Mock response to cache
response_to_cache = {
"id": "test-999",
"choices": [{"message": {"content": "Berlin is the capital of Germany."}}],
}
# Mock the async embedding function
with patch(
"litellm.aembedding",
return_value={"data": [{"embedding": [0.2, 0.2, 0.2]}]},
):
# Test async_set_cache
await qdrant_cache.async_set_cache(
key="test_key",
value=response_to_cache,
messages=[{"content": "What is the capital of Germany?"}],
metadata={},
)
# Verify async upsert was called
qdrant_cache.async_client.put.assert_called()
upsert_payload = qdrant_cache.async_client.put.call_args.kwargs["json"][
"points"
][0]["payload"]
assert (
upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key"
)
def test_qdrant_semantic_cache_custom_vector_size():
"""
Test that QdrantSemanticCache uses a custom vector_size when creating a new collection.
Verifies that the vector size passed to the constructor is used in the Qdrant collection
creation payload instead of the default 1536.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
mock_exists_response = MagicMock()
mock_exists_response.status_code = 200
mock_exists_response.json.return_value = {"result": {"exists": False}}
# Mock the collection creation response
mock_create_response = MagicMock()
mock_create_response.status_code = 200
mock_create_response.json.return_value = {"result": True}
# Mock the collection details response after creation
mock_details_response = MagicMock()
mock_details_response.status_code = 200
mock_details_response.json.return_value = {"result": {"status": "ok"}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.side_effect = [
mock_exists_response,
mock_details_response,
]
mock_sync_client_instance.put.return_value = mock_create_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize with custom vector_size of 768
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection_768",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
vector_size=768,
)
# Verify the vector_size attribute is set correctly
assert qdrant_cache.vector_size == 768
# Verify the PUT call to create the collection used vector_size=768
put_call = next(
call
for call in mock_sync_client_instance.put.call_args_list
if call.kwargs["url"]
== "http://test.qdrant.local/collections/test_collection_768"
)
create_payload = put_call.kwargs["json"]
assert create_payload["vectors"]["size"] == 768
assert create_payload["vectors"]["distance"] == "Cosine"
def test_qdrant_semantic_cache_default_vector_size():
"""
Test that QdrantSemanticCache defaults to QDRANT_VECTOR_SIZE (1536) when vector_size
is not provided, and stores it as self.vector_size.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.return_value = mock_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
from litellm.constants import QDRANT_VECTOR_SIZE
# Initialize without vector_size
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
)
# Verify it falls back to the default QDRANT_VECTOR_SIZE constant
assert qdrant_cache.vector_size == QDRANT_VECTOR_SIZE
def test_qdrant_semantic_cache_large_vector_size():
"""
Test that QdrantSemanticCache supports large embedding dimensions (e.g. 4096, 8192)
for models like Stella, bge-en-icl, etc.
"""
with (
patch(
"litellm.llms.custom_httpx.http_handler._get_httpx_client"
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
mock_exists_response = MagicMock()
mock_exists_response.status_code = 200
mock_exists_response.json.return_value = {"result": {"exists": False}}
mock_create_response = MagicMock()
mock_create_response.status_code = 200
mock_create_response.json.return_value = {"result": True}
mock_details_response = MagicMock()
mock_details_response.status_code = 200
mock_details_response.json.return_value = {"result": {"status": "ok"}}
mock_sync_client_instance = MagicMock()
mock_sync_client_instance.get.side_effect = [
mock_exists_response,
mock_details_response,
]
mock_sync_client_instance.put.return_value = mock_create_response
mock_sync_client.return_value = mock_sync_client_instance
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
# Initialize with a large vector_size of 4096
qdrant_cache = QdrantSemanticCache(
collection_name="test_collection_4096",
qdrant_api_base="http://test.qdrant.local",
qdrant_api_key="test_key",
similarity_threshold=0.8,
vector_size=4096,
)
assert qdrant_cache.vector_size == 4096
# Verify the collection was created with 4096
put_call = next(
call
for call in mock_sync_client_instance.put.call_args_list
if call.kwargs["url"]
== "http://test.qdrant.local/collections/test_collection_4096"
)
create_payload = put_call.kwargs["json"]
assert create_payload["vectors"]["size"] == 4096
def _router_proxy_module(router, model_name):
mod = types.ModuleType("litellm.proxy.proxy_server")
mod.llm_router = router
mod.llm_model_list = [{"model_name": model_name}]
return mod
def test_qdrant_sync_get_cache_routes_through_router(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
cache.qdrant_api_base = "http://test.qdrant.local"
cache.collection_name = "test_collection"
cache.headers = {"Content-Type": "application/json", "api-key": "test_key"}
cache.similarity_threshold = 0.8
cache.sync_client = MagicMock()
search_response = MagicMock()
search_response.status_code = 200
search_response.json.return_value = {"result": []}
cache.sync_client.post.return_value = search_response
router = MagicMock()
router.embedding = MagicMock(
return_value={"data": [{"embedding": [0.3, 0.3, 0.3]}]}
)
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
)
with patch("litellm.embedding") as direct_embed:
result = cache.get_cache(
key="test_key",
messages=[{"content": "What is the capital of France?"}],
metadata={},
)
assert result is None
router.embedding.assert_called_once()
assert router.embedding.call_args.kwargs["model"] == "sem-embed"
direct_embed.assert_not_called()
def test_qdrant_sync_set_cache_falls_back_to_direct(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "text-embedding-ada-002"
cache.qdrant_api_base = "http://test.qdrant.local"
cache.collection_name = "test_collection"
cache.headers = {"Content-Type": "application/json", "api-key": "test_key"}
cache.sync_client = MagicMock()
put_response = MagicMock()
put_response.status_code = 200
cache.sync_client.put.return_value = put_response
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.1, 0.1]}]}
) as direct_embed:
cache.set_cache(
key="test_key",
value={"content": "Paris"},
messages=[{"content": "What is the capital of France?"}],
)
direct_embed.assert_called_once()
@pytest.mark.asyncio
async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
router = MagicMock()
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
)
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"
assert md["semantic-cache-embedding"] is True