mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
1028 lines
37 KiB
Python
1028 lines
37 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
|