litellm/tests/test_litellm/caching/test_qdrant_semantic_cache.py
yassin b5a7032eb4 fix(proxy): run the remaining inline token counts off the event loop
Wrap the context-management editors, the end-of-stream chunk builder,
acount_tokens, the compression interception hook, the passthrough
interrupted-stream recovery, the A2A usage counters, and the semantic
cache embedding truncation in asyncify so a multi-megabyte payload no
longer stalls the worker's event loop while it is tokenized

The pass-through suite now drains the process-global logging worker
from an autouse conftest fixture so work queued on one test's loop
cannot fire against the next test's callbacks

Resolves LIT-7190

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-08 18:42:38 +00:00

1061 lines
38 KiB
Python

import sys
import types
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
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,
embedding_max_input_tokens=512,
)
# 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
assert qdrant_cache.embedding_max_input_tokens == 512
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.get_configured_token_limits.return_value = (None, None)
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.get_configured_token_limits.return_value = (None, None)
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
LONG_PROMPT = " ".join(f"token{i}" for i in range(300))
def _token_count(model, text):
import litellm
return len(litellm.encode(model=model, text=text))
def test_qdrant_get_embedding_truncates_to_deployment_max_input_tokens(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
router = MagicMock()
router.get_configured_token_limits.return_value = (5, None)
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
)
cache._get_embedding(LONG_PROMPT)
sent_input = router.embedding.call_args.kwargs["input"]
assert LONG_PROMPT.startswith(sent_input)
assert _token_count("sem-embed", sent_input) == 5
@pytest.mark.asyncio
async def test_qdrant_async_embedding_explicit_limit_beats_deployment_limit(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
cache.embedding_max_input_tokens = 3
router = MagicMock()
router.get_configured_token_limits.return_value = (8191, None)
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(LONG_PROMPT)
sent_input = router.aembedding.call_args.kwargs["input"]
assert _token_count("sem-embed", sent_input) == 3
@pytest.mark.asyncio
async def test_qdrant_async_embedding_call_is_bounded(monkeypatch):
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
cache.embedding_max_input_tokens = None
cache.embedding_timeout = 1.5
router = MagicMock()
router.get_configured_token_limits.return_value = (None, None)
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("What is the capital of France?")
assert router.aembedding.call_args.kwargs["timeout"] == 1.5
assert router.aembedding.call_args.kwargs["num_retries"] == 0
@pytest.mark.asyncio
async def test_qdrant_async_embedding_gives_up_on_unresponsive_endpoint(monkeypatch):
import asyncio
import time
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
cache.embedding_max_input_tokens = None
cache.embedding_timeout = 0.05
async def never_responds(**kwargs):
await asyncio.sleep(3)
return {"data": [{"embedding": [0.1, 0.2]}]}
router = MagicMock()
router.get_configured_token_limits.return_value = (None, None)
router.aembedding = never_responds
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
)
started = time.monotonic()
with pytest.raises(asyncio.TimeoutError):
await cache._get_async_embedding("What is the capital of France?")
assert time.monotonic() - started < 1.0
def test_qdrant_semantic_cache_defaults_embedding_timeout():
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60
@pytest.mark.asyncio
async def test_qdrant_async_embedding_truncates_off_the_event_loop(monkeypatch):
from tests.large_text import text
from tests.test_litellm.litellm_core_utils.event_loop_lag import (
assert_loop_stayed_free,
timed_with_loop_lags,
warm_tokenizer,
)
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
warm_tokenizer("sem-embed")
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
cache.embedding_model = "sem-embed"
cache.embedding_max_input_tokens = 5
cache.embedding_timeout = 5
router = MagicMock()
router.get_configured_token_limits.return_value = (8191, None)
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"),
)
response, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100))
assert response["data"][0]["embedding"] == [0.1, 0.2]
assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5
assert_loop_stayed_free(took, lags)