mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +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
370 lines
13 KiB
Python
370 lines
13 KiB
Python
"""
|
|
Test suite for router embedding method header propagation.
|
|
|
|
This tests the fix for the issue where the embedding method was not
|
|
propagating proxy model configuration headers to the LLM API calls.
|
|
|
|
The fix ensures that router.embedding() calls _update_kwargs_before_fallbacks()
|
|
just like router.completion() does, which properly sets up metadata and allows
|
|
default_litellm_params (including headers) to be propagated.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, patch, AsyncMock
|
|
|
|
import pytest
|
|
|
|
|
|
from litellm import Router
|
|
|
|
|
|
class TestRouterEmbeddingHeaders:
|
|
"""Test that embedding methods properly propagate headers from router configuration."""
|
|
|
|
def test_embedding_calls_update_kwargs_before_fallbacks(self):
|
|
"""
|
|
Test that router.embedding() calls _update_kwargs_before_fallbacks.
|
|
|
|
This ensures that metadata is properly set up before the fallback mechanism,
|
|
which is necessary for header propagation to work correctly.
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock the _update_kwargs_before_fallbacks method to verify it's called
|
|
with patch.object(
|
|
router,
|
|
"_update_kwargs_before_fallbacks",
|
|
wraps=router._update_kwargs_before_fallbacks,
|
|
) as mock_update:
|
|
with patch("litellm.embedding") as mock_litellm_embedding:
|
|
mock_litellm_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="text-embedding-3-small", input=["test input"])
|
|
|
|
# Verify _update_kwargs_before_fallbacks was called
|
|
mock_update.assert_called_once()
|
|
call_kwargs = mock_update.call_args[1]
|
|
assert call_kwargs["model"] == "text-embedding-3-small"
|
|
assert "kwargs" in call_kwargs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aembedding_calls_update_kwargs_before_fallbacks(self):
|
|
"""
|
|
Test that router.aembedding() calls _update_kwargs_before_fallbacks.
|
|
|
|
This ensures consistency between sync and async embedding methods.
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
# Mock the _update_kwargs_before_fallbacks method to verify it's called
|
|
with patch.object(
|
|
router,
|
|
"_update_kwargs_before_fallbacks",
|
|
wraps=router._update_kwargs_before_fallbacks,
|
|
) as mock_update:
|
|
with patch(
|
|
"litellm.aembedding", new_callable=AsyncMock
|
|
) as mock_litellm_aembedding:
|
|
mock_litellm_aembedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
await router.aembedding(
|
|
model="text-embedding-3-small", input=["test input"]
|
|
)
|
|
|
|
# Verify _update_kwargs_before_fallbacks was called
|
|
mock_update.assert_called_once()
|
|
call_kwargs = mock_update.call_args[1]
|
|
assert call_kwargs["model"] == "text-embedding-3-small"
|
|
assert "kwargs" in call_kwargs
|
|
|
|
def test_embedding_propagates_default_litellm_params(self):
|
|
"""
|
|
Test that embedding calls properly propagate default_litellm_params including headers.
|
|
|
|
This is the main fix - ensuring that headers set in default_litellm_params
|
|
are included in the embedding request.
|
|
"""
|
|
custom_headers = {"X-Custom-Header": "test-value", "X-API-Version": "v2"}
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
# Create router with default_litellm_params containing headers
|
|
router = Router(
|
|
model_list=model_list,
|
|
default_litellm_params={
|
|
"headers": custom_headers,
|
|
"metadata": {"test_key": "test_value"},
|
|
},
|
|
)
|
|
|
|
with patch("litellm.embedding") as mock_litellm_embedding:
|
|
mock_litellm_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="text-embedding-3-small", input=["test input"])
|
|
|
|
# Verify that litellm.embedding was called with the headers
|
|
mock_litellm_embedding.assert_called_once()
|
|
call_kwargs = mock_litellm_embedding.call_args[1]
|
|
|
|
# Check that headers were included
|
|
assert "headers" in call_kwargs
|
|
assert call_kwargs["headers"] == custom_headers
|
|
|
|
# Check that metadata was properly set up
|
|
assert "metadata" in call_kwargs
|
|
assert "model_group" in call_kwargs["metadata"]
|
|
assert call_kwargs["metadata"]["model_group"] == "text-embedding-3-small"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aembedding_propagates_default_litellm_params(self):
|
|
"""
|
|
Test that async embedding calls properly propagate default_litellm_params including headers.
|
|
"""
|
|
custom_headers = {"X-Custom-Header": "test-value", "X-API-Version": "v2"}
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
# Create router with default_litellm_params containing headers
|
|
router = Router(
|
|
model_list=model_list,
|
|
default_litellm_params={
|
|
"headers": custom_headers,
|
|
"metadata": {"test_key": "test_value"},
|
|
},
|
|
)
|
|
|
|
with patch(
|
|
"litellm.aembedding", new_callable=AsyncMock
|
|
) as mock_litellm_aembedding:
|
|
mock_litellm_aembedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
await router.aembedding(
|
|
model="text-embedding-3-small", input=["test input"]
|
|
)
|
|
|
|
# Verify that litellm.aembedding was called with the headers
|
|
mock_litellm_aembedding.assert_called_once()
|
|
call_kwargs = mock_litellm_aembedding.call_args[1]
|
|
|
|
# Check that headers were included
|
|
assert "headers" in call_kwargs
|
|
assert call_kwargs["headers"] == custom_headers
|
|
|
|
# Check that metadata was properly set up
|
|
assert "metadata" in call_kwargs
|
|
assert "model_group" in call_kwargs["metadata"]
|
|
assert call_kwargs["metadata"]["model_group"] == "text-embedding-3-small"
|
|
|
|
def test_embedding_metadata_includes_model_group(self):
|
|
"""
|
|
Test that embedding calls include model_group in metadata.
|
|
|
|
The _update_kwargs_before_fallbacks method should set this up.
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "test-embedding-model",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
with patch("litellm.embedding") as mock_litellm_embedding:
|
|
mock_litellm_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="test-embedding-model", input=["test input"])
|
|
|
|
call_kwargs = mock_litellm_embedding.call_args[1]
|
|
|
|
# Verify metadata contains model_group
|
|
assert "metadata" in call_kwargs
|
|
assert "model_group" in call_kwargs["metadata"]
|
|
assert call_kwargs["metadata"]["model_group"] == "test-embedding-model"
|
|
|
|
def test_embedding_sets_num_retries_from_router(self):
|
|
"""
|
|
Test that embedding calls inherit num_retries from router configuration.
|
|
|
|
This is set by _update_kwargs_before_fallbacks.
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
# Create router with num_retries set
|
|
router = Router(model_list=model_list, num_retries=3)
|
|
|
|
with patch("litellm.embedding") as mock_litellm_embedding:
|
|
mock_litellm_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="text-embedding-3-small", input=["test input"])
|
|
|
|
# Verify num_retries was not set in the call (it's handled by function_with_fallbacks)
|
|
# The important thing is that it was set in kwargs before being passed to function_with_fallbacks
|
|
# We verify this indirectly by checking that _update_kwargs_before_fallbacks was called
|
|
mock_litellm_embedding.assert_called_once()
|
|
|
|
def test_embedding_sets_litellm_trace_id(self):
|
|
"""
|
|
Test that embedding calls include a litellm_trace_id.
|
|
|
|
This is generated and set by _update_kwargs_before_fallbacks.
|
|
"""
|
|
model_list = [
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
}
|
|
]
|
|
|
|
router = Router(model_list=model_list)
|
|
|
|
with patch("litellm.embedding") as mock_litellm_embedding:
|
|
mock_litellm_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="text-embedding-3-small", input=["test input"])
|
|
|
|
call_kwargs = mock_litellm_embedding.call_args[1]
|
|
|
|
# Verify litellm_trace_id was set
|
|
assert "litellm_trace_id" in call_kwargs
|
|
assert isinstance(call_kwargs["litellm_trace_id"], str)
|
|
assert len(call_kwargs["litellm_trace_id"]) > 0
|
|
|
|
def test_embedding_consistency_with_completion(self):
|
|
"""
|
|
Test that embedding and completion methods handle kwargs similarly.
|
|
|
|
Both should call _update_kwargs_before_fallbacks to ensure consistent behavior.
|
|
"""
|
|
custom_headers = {"X-Test": "value"}
|
|
|
|
model_list = [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": "fake-key",
|
|
},
|
|
},
|
|
{
|
|
"model_name": "text-embedding-3-small",
|
|
"litellm_params": {
|
|
"model": "text-embedding-3-small",
|
|
"api_key": "fake-key",
|
|
},
|
|
},
|
|
]
|
|
|
|
router = Router(
|
|
model_list=model_list, default_litellm_params={"headers": custom_headers}
|
|
)
|
|
|
|
# Test completion
|
|
with patch("litellm.completion") as mock_completion:
|
|
mock_completion.return_value = MagicMock()
|
|
|
|
router.completion(
|
|
model="gpt-5-mini", messages=[{"role": "user", "content": "test"}]
|
|
)
|
|
|
|
completion_kwargs = mock_completion.call_args[1]
|
|
|
|
# Test embedding
|
|
with patch("litellm.embedding") as mock_embedding:
|
|
mock_embedding.return_value = MagicMock(
|
|
data=[{"embedding": [0.1, 0.2, 0.3]}]
|
|
)
|
|
|
|
router.embedding(model="text-embedding-3-small", input=["test input"])
|
|
|
|
embedding_kwargs = mock_embedding.call_args[1]
|
|
|
|
# Both should have headers from default_litellm_params
|
|
assert "headers" in completion_kwargs
|
|
assert "headers" in embedding_kwargs
|
|
assert completion_kwargs["headers"] == custom_headers
|
|
assert embedding_kwargs["headers"] == custom_headers
|
|
|
|
# Both should have metadata with model_group
|
|
assert "metadata" in completion_kwargs
|
|
assert "metadata" in embedding_kwargs
|
|
assert "model_group" in completion_kwargs["metadata"]
|
|
assert "model_group" in embedding_kwargs["metadata"]
|
|
|
|
# Both should have litellm_trace_id
|
|
assert "litellm_trace_id" in completion_kwargs
|
|
assert "litellm_trace_id" in embedding_kwargs
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run a simple test
|
|
test = TestRouterEmbeddingHeaders()
|
|
test.test_embedding_calls_update_kwargs_before_fallbacks()
|
|
test.test_embedding_propagates_default_litellm_params()
|
|
test.test_embedding_metadata_includes_model_group()
|
|
test.test_embedding_sets_litellm_trace_id()
|
|
test.test_embedding_consistency_with_completion()
|
|
print("All tests passed!") # noqa: T201
|