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>
910 lines
33 KiB
Python
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
|