mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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>
1061 lines
38 KiB
Python
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)
|