mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 02:24:54 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix-config-update-targeted-upserts
This commit is contained in:
@@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive(
|
||||
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
|
||||
engine_client,
|
||||
) -> None:
|
||||
"""_run_reconnect_cycle uses disconnect/connect when engine is alive."""
|
||||
engine_client._engine_pid = 1234
|
||||
"""Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1.
|
||||
|
||||
with patch.object(engine_client, "_is_engine_alive", return_value=True):
|
||||
The old "lightweight" path called `disconnect()` + `connect()`, which
|
||||
blocks the event loop on the sync `process.wait()` inside aclose().
|
||||
The fix routes both engine-alive and engine-dead paths through
|
||||
`recreate_prisma_client`, which non-blockingly kills the old engine.
|
||||
"""
|
||||
engine_client._engine_pid = 1234
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(engine_client, "_is_engine_alive", return_value=True),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
engine_client.db.connect.assert_awaited_once()
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
|
||||
"postgresql://test"
|
||||
)
|
||||
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
engine_client.db.recreate_prisma_client.assert_not_awaited()
|
||||
engine_client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown(
|
||||
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
|
||||
engine_client,
|
||||
) -> None:
|
||||
"""_run_reconnect_cycle uses lightweight path when engine PID is not tracked."""
|
||||
"""When the engine PID is not tracked, direct reconnect still runs."""
|
||||
engine_client._engine_pid = 0
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
engine_client.db.connect.assert_awaited_once()
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
|
||||
"postgresql://test"
|
||||
)
|
||||
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
engine_client.db.recreate_prisma_client.assert_not_awaited()
|
||||
engine_client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_escalation_after_consecutive_lightweight_failures(engine_client):
|
||||
"""After N consecutive lightweight reconnect failures, _engine_confirmed_dead
|
||||
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
|
||||
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
|
||||
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
|
||||
engine_client._reconnect_escalation_threshold = 3
|
||||
engine_client._consecutive_reconnect_failures = 0
|
||||
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
|
||||
engine_client._start_engine_watcher = AsyncMock(return_value=None)
|
||||
|
||||
# Make lightweight reconnect fail every time
|
||||
engine_client.db.disconnect = AsyncMock(return_value=None)
|
||||
engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed"))
|
||||
# Make direct reconnect fail every time
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=Exception("recreate failed")
|
||||
)
|
||||
|
||||
# Run 3 failed reconnect attempts
|
||||
for i in range(3):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
assert result is False
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
for _ in range(3):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
assert result is False
|
||||
|
||||
assert engine_client._consecutive_reconnect_failures == 3
|
||||
|
||||
# Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle
|
||||
# Next attempt should escalate to the heavy path (recreate_prisma_client still
|
||||
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
engine_client._start_engine_watcher = AsyncMock(return_value=None)
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test_escalation", timeout_seconds=5.0
|
||||
)
|
||||
|
||||
# Heavy reconnect should have been attempted (recreate_prisma_client called)
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once()
|
||||
|
||||
|
||||
@@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client):
|
||||
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
|
||||
engine_client._consecutive_reconnect_failures = 2
|
||||
engine_client._db_reconnect_cooldown_seconds = 0
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
# Make reconnect succeed
|
||||
engine_client.db.disconnect = AsyncMock(return_value=None)
|
||||
engine_client.db.connect = AsyncMock(return_value=None)
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
|
||||
assert result is True
|
||||
assert engine_client._consecutive_reconnect_failures == 0
|
||||
|
||||
@@ -314,11 +314,11 @@ def test_update_litellm_params_for_health_check():
|
||||
# Issue #15807: Fixes health checks sending "region/model" as model ID to AWS
|
||||
model_info = {}
|
||||
litellm_params = {
|
||||
"model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"model": "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"api_key": "fake_key",
|
||||
}
|
||||
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
|
||||
assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
assert updated_params["model"] == "anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
|
||||
# Test with Bedrock cross-region inference profile - should preserve the inference profile prefix
|
||||
# AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing
|
||||
|
||||
@@ -366,3 +366,40 @@ def test_generic_api_compatible_callbacks_json_unknown_callback():
|
||||
# Should return the string unchanged
|
||||
assert result == "unknown_callback", "Unknown callback should be returned as-is"
|
||||
assert isinstance(result, str), "Unknown callback should remain a string"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_settings_retry_config():
|
||||
"""
|
||||
Test that generic_api callback_settings are passed to GenericAPILogger.
|
||||
"""
|
||||
from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger
|
||||
from litellm.litellm_core_utils.logging_callback_manager import (
|
||||
_generic_api_logger_cache,
|
||||
)
|
||||
|
||||
callback_name = "test_generic_api_retry_config"
|
||||
_generic_api_logger_cache.pop(callback_name, None)
|
||||
litellm.callback_settings[callback_name] = {
|
||||
"callback_type": "generic_api",
|
||||
"endpoint": "https://example.com/api/logs",
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"max_retries": 2,
|
||||
"retry_delay": 0.5,
|
||||
"timeout": 3,
|
||||
}
|
||||
|
||||
try:
|
||||
result = LoggingCallbackManager._add_custom_callback_generic_api_str(
|
||||
callback_name
|
||||
)
|
||||
|
||||
assert isinstance(result, GenericAPILogger)
|
||||
assert result.endpoint == "https://example.com/api/logs"
|
||||
assert result.headers == {"Content-Type": "application/json"}
|
||||
assert result.max_retries == 2
|
||||
assert result.retry_delay == 0.5
|
||||
assert result.timeout == 3
|
||||
finally:
|
||||
litellm.callback_settings.pop(callback_name, None)
|
||||
_generic_api_logger_cache.pop(callback_name, None)
|
||||
|
||||
@@ -2309,11 +2309,11 @@ def test_get_provider_audio_transcription_config():
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected_bool",
|
||||
[
|
||||
("anthropic.claude-3-7-sonnet-20250219-v1:0", True),
|
||||
("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True),
|
||||
("anthropic.claude-sonnet-4-5-20250929-v1:0", True),
|
||||
("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True),
|
||||
],
|
||||
)
|
||||
def test_claude_3_7_sonnet_supports_pdf_input(model, expected_bool):
|
||||
def test_claude_sonnet_4_5_supports_pdf_input(model, expected_bool):
|
||||
from litellm.utils import supports_pdf_input
|
||||
|
||||
assert supports_pdf_input(model) == expected_bool
|
||||
|
||||
@@ -134,7 +134,7 @@ class TestBedrockAnthropicPromptCachingRegression:
|
||||
if "converse" in model_prefix:
|
||||
config = AmazonConverseConfig()
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
@@ -162,7 +162,7 @@ class TestBedrockAnthropicPromptCachingRegression:
|
||||
else:
|
||||
config = AmazonAnthropicClaudeConfig()
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
@@ -227,7 +227,7 @@ class TestBedrockAnthropicPromptCachingRegression:
|
||||
if "converse" in model_prefix:
|
||||
config = AmazonConverseConfig()
|
||||
result = config._transform_request_helper(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
system_content_blocks=[],
|
||||
optional_params={},
|
||||
messages=messages,
|
||||
@@ -236,7 +236,7 @@ class TestBedrockAnthropicPromptCachingRegression:
|
||||
else:
|
||||
config = AmazonAnthropicClaudeConfig()
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
@@ -498,7 +498,7 @@ class TestBedrockAnthropicCombinedRegressions:
|
||||
if "converse" in model_prefix:
|
||||
config = AmazonConverseConfig()
|
||||
result = config._transform_request_helper(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
system_content_blocks=[],
|
||||
optional_params={},
|
||||
messages=messages,
|
||||
@@ -518,7 +518,7 @@ class TestBedrockAnthropicCombinedRegressions:
|
||||
else:
|
||||
config = AmazonAnthropicClaudeConfig()
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
|
||||
@@ -1323,7 +1323,7 @@ def test_base_aws_llm_get_credentials():
|
||||
def test_bedrock_completion_test_2():
|
||||
litellm.set_verbose = True
|
||||
data = {
|
||||
"model": "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
@@ -1630,7 +1630,7 @@ def test_bedrock_completion_test_4(modify_params):
|
||||
litellm.modify_params = modify_params
|
||||
|
||||
data = {
|
||||
"model": "anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"model": "anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
@@ -2115,7 +2115,7 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest):
|
||||
|
||||
def get_base_completion_call_args_with_thinking(self) -> dict:
|
||||
return {
|
||||
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"thinking": {"type": "enabled", "budget_tokens": 16000},
|
||||
}
|
||||
|
||||
@@ -2828,7 +2828,7 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode):
|
||||
client = AsyncHTTPHandler()
|
||||
|
||||
params = {
|
||||
"model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
@@ -2887,7 +2887,7 @@ async def test_bedrock_stream_thinking_content_openwebui():
|
||||
```
|
||||
"""
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[{"role": "user", "content": "Hello who is this?"}],
|
||||
stream=True,
|
||||
max_tokens=1080,
|
||||
|
||||
@@ -580,7 +580,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers():
|
||||
def test_litellm_gateway_from_sdk_with_thinking_param():
|
||||
try:
|
||||
response = litellm.completion(
|
||||
model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[{"role": "user", "content": "Hello world"}],
|
||||
api_base="http://0.0.0.0:4000",
|
||||
api_key="sk-PIp1h0RekR",
|
||||
|
||||
@@ -1828,7 +1828,7 @@ def test_azure_response_format_param():
|
||||
"model, provider",
|
||||
[
|
||||
("claude-3-7-sonnet-20240620-v1:0", "anthropic"),
|
||||
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
|
||||
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
|
||||
("invoke/anthropic.claude-3-7-sonnet-20240620-v1:0", "bedrock"),
|
||||
("claude-3-7-sonnet@20250219", "vertex_ai"),
|
||||
],
|
||||
|
||||
@@ -3493,8 +3493,14 @@ def test_litellm_api_base(monkeypatch, provider, route):
|
||||
|
||||
|
||||
def test_gemini_tool_calling_working_demo():
|
||||
load_vertex_ai_credentials()
|
||||
litellm._turn_on_debug()
|
||||
"""
|
||||
Regression test: tool params with anyOf containing a `{"type": "array"}`
|
||||
branch (no items field at all) must synthesize items before the request
|
||||
is sent to Vertex (Vertex rejects array types missing items).
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
args = {
|
||||
"messages": [
|
||||
{
|
||||
@@ -3564,13 +3570,75 @@ def test_gemini_tool_calling_working_demo():
|
||||
],
|
||||
"vertex_location": "global",
|
||||
}
|
||||
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
|
||||
print(response)
|
||||
|
||||
client = HTTPHandler()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"Content-Type": "application/json"}
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{"text": "Hello!"}],
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 5,
|
||||
"totalTokenCount": 15,
|
||||
},
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(client, "post", return_value=mock_response) as mock_post,
|
||||
patch.object(
|
||||
VertexBase,
|
||||
"_ensure_access_token",
|
||||
return_value=("fake-token", "fake-project"),
|
||||
),
|
||||
):
|
||||
completion(
|
||||
model="vertex_ai/gemini-3-flash-preview",
|
||||
client=client,
|
||||
**args,
|
||||
)
|
||||
|
||||
sent_body = mock_post.call_args.kwargs.get(
|
||||
"json"
|
||||
) or mock_post.call_args.kwargs.get("data")
|
||||
assert sent_body is not None, "expected request body to be sent"
|
||||
if isinstance(sent_body, str):
|
||||
sent_body = json.loads(sent_body)
|
||||
|
||||
function_decl = sent_body["tools"][0]["function_declarations"][0]
|
||||
callbacks_schema = function_decl["parameters"]["properties"]["config"][
|
||||
"properties"
|
||||
]["callbacks"]
|
||||
array_branches = [
|
||||
branch
|
||||
for branch in callbacks_schema["anyOf"]
|
||||
if branch.get("type", "").lower() == "array"
|
||||
]
|
||||
assert array_branches, "expected an array branch in callbacks anyOf"
|
||||
for branch in array_branches:
|
||||
assert "items" in branch and branch["items"], (
|
||||
f"array branch in callbacks.anyOf must include non-empty items "
|
||||
f"(Vertex rejects array types missing items). Got: {branch}"
|
||||
)
|
||||
|
||||
|
||||
def test_gemini_tool_calling_not_working():
|
||||
load_vertex_ai_credentials()
|
||||
litellm._turn_on_debug()
|
||||
"""
|
||||
Regression test: tool params with anyOf containing both an empty-items
|
||||
array branch and a null branch must serialize with items present on the
|
||||
array branch (Vertex rejects array types missing `items`).
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
args = {
|
||||
"messages": [
|
||||
{
|
||||
@@ -3637,8 +3705,64 @@ def test_gemini_tool_calling_not_working():
|
||||
],
|
||||
"vertex_location": "global",
|
||||
}
|
||||
response = completion(model="vertex_ai/gemini-3-flash-preview", **args)
|
||||
print(response)
|
||||
|
||||
client = HTTPHandler()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"Content-Type": "application/json"}
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{"text": "Hello!"}],
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 5,
|
||||
"totalTokenCount": 15,
|
||||
},
|
||||
}
|
||||
|
||||
with (
|
||||
patch.object(client, "post", return_value=mock_response) as mock_post,
|
||||
patch.object(
|
||||
VertexBase,
|
||||
"_ensure_access_token",
|
||||
return_value=("fake-token", "fake-project"),
|
||||
),
|
||||
):
|
||||
completion(
|
||||
model="vertex_ai/gemini-3-flash-preview",
|
||||
client=client,
|
||||
**args,
|
||||
)
|
||||
|
||||
sent_body = mock_post.call_args.kwargs.get(
|
||||
"json"
|
||||
) or mock_post.call_args.kwargs.get("data")
|
||||
assert sent_body is not None, "expected request body to be sent"
|
||||
if isinstance(sent_body, str):
|
||||
sent_body = json.loads(sent_body)
|
||||
|
||||
function_decl = sent_body["tools"][0]["function_declarations"][0]
|
||||
callbacks_schema = function_decl["parameters"]["properties"]["config"][
|
||||
"properties"
|
||||
]["callbacks"]
|
||||
array_branches = [
|
||||
branch
|
||||
for branch in callbacks_schema["anyOf"]
|
||||
if branch.get("type", "").lower() == "array"
|
||||
]
|
||||
assert array_branches, "expected an array branch in callbacks anyOf"
|
||||
for branch in array_branches:
|
||||
assert "items" in branch and branch["items"], (
|
||||
f"array branch in callbacks.anyOf must include non-empty items "
|
||||
f"(Vertex rejects array types missing items). Got: {branch}"
|
||||
)
|
||||
|
||||
|
||||
def test_vertex_ai_llama_tool_calling():
|
||||
|
||||
@@ -159,7 +159,7 @@ def test_aaparallel_function_call(model):
|
||||
"model",
|
||||
[
|
||||
"anthropic/claude-4-sonnet-20250514",
|
||||
"bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
],
|
||||
)
|
||||
@pytest.mark.flaky(retries=3, delay=1)
|
||||
|
||||
@@ -30,10 +30,9 @@ def test_model_alias_map(caplog):
|
||||
)
|
||||
print(response.model)
|
||||
|
||||
captured_logs = [rec.levelname for rec in caplog.records]
|
||||
|
||||
for log in captured_logs:
|
||||
assert "ERROR" not in log
|
||||
for rec in caplog.records:
|
||||
if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"):
|
||||
pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}")
|
||||
|
||||
assert "llama-3.1-8b-instant" in response.model
|
||||
except litellm.ServiceUnavailableError:
|
||||
|
||||
@@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
|
||||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=litellm_params_dict,
|
||||
extra_body=None,
|
||||
)
|
||||
)
|
||||
captured_request_body["url"] = url
|
||||
|
||||
@@ -8,6 +8,7 @@ sys.path.insert(0, os.path.abspath("../.."))
|
||||
import asyncio
|
||||
import litellm
|
||||
import gzip
|
||||
import httpx
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
@@ -470,3 +471,96 @@ async def test_generic_api_callback_invalid_log_format():
|
||||
endpoint=test_endpoint,
|
||||
log_format="invalid_format", # type: ignore # Intentionally invalid for testing
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_retries_timeout_then_succeeds():
|
||||
"""
|
||||
Test that GenericAPILogger retries LiteLLM timeout errors when configured.
|
||||
"""
|
||||
test_endpoint = "https://example.com/api/logs"
|
||||
generic_logger = GenericAPILogger(
|
||||
endpoint=test_endpoint,
|
||||
max_retries=1,
|
||||
retry_delay=0,
|
||||
timeout=0.2,
|
||||
)
|
||||
|
||||
mock_post = AsyncMock()
|
||||
mock_post.side_effect = [
|
||||
litellm.Timeout(
|
||||
message="Connection timed out",
|
||||
model="default-model-name",
|
||||
llm_provider="litellm-httpx-handler",
|
||||
),
|
||||
type("Response", (), {"status_code": 200})(),
|
||||
]
|
||||
generic_logger.async_httpx_client.post = mock_post
|
||||
generic_logger.log_queue = [{"event": "timeout-retry"}]
|
||||
|
||||
await generic_logger.async_send_batch()
|
||||
|
||||
assert mock_post.call_count == 2
|
||||
first_call = mock_post.call_args_list[0][1]
|
||||
assert first_call["url"] == test_endpoint
|
||||
assert first_call["timeout"] == 0.2
|
||||
assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_retries_5xx_then_succeeds():
|
||||
"""
|
||||
Test that GenericAPILogger retries transient HTTP 5xx errors when configured.
|
||||
"""
|
||||
test_endpoint = "https://example.com/api/logs"
|
||||
generic_logger = GenericAPILogger(
|
||||
endpoint=test_endpoint,
|
||||
max_retries=1,
|
||||
retry_delay=0,
|
||||
)
|
||||
|
||||
request = httpx.Request("POST", test_endpoint)
|
||||
response = httpx.Response(status_code=503, request=request)
|
||||
mock_post = AsyncMock()
|
||||
mock_post.side_effect = [
|
||||
httpx.HTTPStatusError(
|
||||
"Server error",
|
||||
request=request,
|
||||
response=response,
|
||||
),
|
||||
type("Response", (), {"status_code": 200})(),
|
||||
]
|
||||
generic_logger.async_httpx_client.post = mock_post
|
||||
generic_logger.log_queue = [{"event": "5xx-retry"}]
|
||||
|
||||
await generic_logger.async_send_batch()
|
||||
|
||||
assert mock_post.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generic_api_callback_does_not_retry_4xx():
|
||||
"""
|
||||
Test that GenericAPILogger does not retry non-transient HTTP 4xx errors.
|
||||
"""
|
||||
test_endpoint = "https://example.com/api/logs"
|
||||
generic_logger = GenericAPILogger(
|
||||
endpoint=test_endpoint,
|
||||
max_retries=2,
|
||||
retry_delay=0,
|
||||
)
|
||||
|
||||
request = httpx.Request("POST", test_endpoint)
|
||||
response = httpx.Response(status_code=401, request=request)
|
||||
mock_post = AsyncMock()
|
||||
mock_post.side_effect = httpx.HTTPStatusError(
|
||||
"Unauthorized",
|
||||
request=request,
|
||||
response=response,
|
||||
)
|
||||
generic_logger.async_httpx_client.post = mock_post
|
||||
generic_logger.log_queue = [{"event": "4xx-no-retry"}]
|
||||
|
||||
await generic_logger.async_send_batch()
|
||||
|
||||
mock_post.assert_called_once()
|
||||
|
||||
@@ -96,8 +96,8 @@ class BaseAnthropicMessagesPromptCachingTest(ABC):
|
||||
Returns the model string to use for tests.
|
||||
|
||||
Examples:
|
||||
- "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
- "bedrock/invoke/anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
- "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
- "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest):
|
||||
"""
|
||||
|
||||
def get_model(self) -> str:
|
||||
return "bedrock/converse/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
|
||||
|
||||
class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
|
||||
@@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest):
|
||||
"""
|
||||
|
||||
def get_model(self) -> str:
|
||||
return "bedrock/invoke/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
|
||||
@@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response():
|
||||
|
||||
print(f"response: {response}")
|
||||
assert len(response.data) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_cache_preserves_prompt_tokens_details():
|
||||
"""Test that prompt_tokens_details (including image_count) survives a full cache hit."""
|
||||
llm_caching_handler = LLMCachingHandler(
|
||||
original_function=MagicMock(),
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
cached_result = [
|
||||
{
|
||||
"embedding": [-0.025, -0.019],
|
||||
"index": 0,
|
||||
"object": "embedding",
|
||||
"model": "amazon.titan-embed-image-v1",
|
||||
"prompt_tokens_details": {"image_count": 1},
|
||||
}
|
||||
]
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
|
||||
final_embedding_cached_response=None,
|
||||
cached_result=cached_result,
|
||||
kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"},
|
||||
logging_obj=mock_logging_obj,
|
||||
start_time=datetime.now(),
|
||||
model="amazon.titan-embed-image-v1",
|
||||
)
|
||||
|
||||
assert cache_hit
|
||||
assert response.usage is not None
|
||||
assert response.usage.prompt_tokens_details is not None
|
||||
assert response.usage.prompt_tokens_details.image_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_cache_backward_compat_no_prompt_tokens_details():
|
||||
"""Test that old cached items without prompt_tokens_details still work."""
|
||||
llm_caching_handler = LLMCachingHandler(
|
||||
original_function=MagicMock(),
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Old-format cached item — no prompt_tokens_details field
|
||||
cached_result = [
|
||||
{
|
||||
"embedding": [-0.025, -0.019],
|
||||
"index": 0,
|
||||
"object": "embedding",
|
||||
"model": "text-embedding-ada-002",
|
||||
}
|
||||
]
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
|
||||
final_embedding_cached_response=None,
|
||||
cached_result=cached_result,
|
||||
kwargs={"model": "text-embedding-ada-002", "input": "test"},
|
||||
logging_obj=mock_logging_obj,
|
||||
start_time=datetime.now(),
|
||||
model="text-embedding-ada-002",
|
||||
)
|
||||
|
||||
assert cache_hit
|
||||
assert response.usage is not None
|
||||
assert response.usage.prompt_tokens_details is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embedding_cache_aggregates_multiple_image_counts():
|
||||
"""Test that image_count is summed correctly across multiple cached items."""
|
||||
llm_caching_handler = LLMCachingHandler(
|
||||
original_function=MagicMock(),
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
cached_result = [
|
||||
{
|
||||
"embedding": [-0.025, -0.019],
|
||||
"index": 0,
|
||||
"object": "embedding",
|
||||
"model": "amazon.titan-embed-image-v1",
|
||||
"prompt_tokens_details": {"image_count": 1},
|
||||
},
|
||||
{
|
||||
"embedding": [0.031, 0.042],
|
||||
"index": 1,
|
||||
"object": "embedding",
|
||||
"model": "amazon.titan-embed-image-v1",
|
||||
"prompt_tokens_details": {"image_count": 1},
|
||||
},
|
||||
]
|
||||
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.async_success_handler = AsyncMock()
|
||||
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
|
||||
final_embedding_cached_response=None,
|
||||
cached_result=cached_result,
|
||||
kwargs={
|
||||
"model": "amazon.titan-embed-image-v1",
|
||||
"input": ["img1", "img2"],
|
||||
},
|
||||
logging_obj=mock_logging_obj,
|
||||
start_time=datetime.now(),
|
||||
model="amazon.titan-embed-image-v1",
|
||||
)
|
||||
|
||||
assert cache_hit
|
||||
assert response.usage.prompt_tokens_details is not None
|
||||
assert response.usage.prompt_tokens_details.image_count == 2
|
||||
|
||||
|
||||
def test_combine_usage_merges_prompt_tokens_details():
|
||||
"""Test that combine_usage merges prompt_tokens_details from both Usage objects."""
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
|
||||
|
||||
llm_caching_handler = LLMCachingHandler(
|
||||
original_function=MagicMock(),
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
usage1 = Usage(
|
||||
prompt_tokens=10,
|
||||
completion_tokens=0,
|
||||
total_tokens=10,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
|
||||
)
|
||||
usage2 = Usage(
|
||||
prompt_tokens=20,
|
||||
completion_tokens=0,
|
||||
total_tokens=20,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2),
|
||||
)
|
||||
|
||||
combined = llm_caching_handler.combine_usage(usage1, usage2)
|
||||
|
||||
assert combined.prompt_tokens == 30
|
||||
assert combined.total_tokens == 30
|
||||
assert combined.prompt_tokens_details is not None
|
||||
assert combined.prompt_tokens_details.image_count == 3
|
||||
|
||||
|
||||
def test_combine_usage_handles_none_details():
|
||||
"""Test that combine_usage works when one or both sides have null prompt_tokens_details."""
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
|
||||
|
||||
llm_caching_handler = LLMCachingHandler(
|
||||
original_function=MagicMock(),
|
||||
request_kwargs={},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
|
||||
# Both null
|
||||
usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10)
|
||||
usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20)
|
||||
combined = llm_caching_handler.combine_usage(usage_a, usage_b)
|
||||
assert combined.prompt_tokens_details is None
|
||||
|
||||
# Only first has details
|
||||
usage_c = Usage(
|
||||
prompt_tokens=10,
|
||||
completion_tokens=0,
|
||||
total_tokens=10,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1),
|
||||
)
|
||||
combined = llm_caching_handler.combine_usage(usage_c, usage_b)
|
||||
assert combined.prompt_tokens_details is not None
|
||||
assert combined.prompt_tokens_details.image_count == 1
|
||||
|
||||
# Only second has details
|
||||
combined = llm_caching_handler.combine_usage(usage_a, usage_c)
|
||||
assert combined.prompt_tokens_details is not None
|
||||
assert combined.prompt_tokens_details.image_count == 1
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1 +1 @@
|
||||
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
|
||||
{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}
|
||||
@@ -220,7 +220,7 @@ async def test_anthropic_cache_control_hook_negative_indices():
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
# Test with multiple messages and negative indices
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
@@ -352,7 +352,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds_logging():
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
cache_control_injection_points=[
|
||||
{"location": "message", "index": 10}
|
||||
@@ -420,7 +420,7 @@ async def test_anthropic_cache_control_hook_negative_out_of_bounds_logging():
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
cache_control_injection_points=[
|
||||
{
|
||||
@@ -486,7 +486,7 @@ async def test_anthropic_cache_control_hook_multiple_user_messages():
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
# Test with multiple user messages and negative indices
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
@@ -586,7 +586,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds(bad_index):
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
cache_control_injection_points=[
|
||||
{"location": "message", "index": bad_index}
|
||||
@@ -651,7 +651,7 @@ async def test_anthropic_cache_control_hook_single_message(message_list):
|
||||
client = AsyncHTTPHandler()
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=message_list,
|
||||
cache_control_injection_points=[{"location": "message", "index": -1}],
|
||||
client=client,
|
||||
@@ -691,7 +691,7 @@ async def test_anthropic_cache_control_hook_empty_message_list():
|
||||
match="bedrock requires at least one non-system message",
|
||||
):
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[],
|
||||
cache_control_injection_points=[
|
||||
{"location": "message", "index": -1}
|
||||
@@ -742,7 +742,7 @@ async def test_anthropic_cache_control_hook_no_op():
|
||||
]
|
||||
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=messages,
|
||||
# No cache_control_injection_points parameter
|
||||
client=client,
|
||||
@@ -799,7 +799,7 @@ async def test_anthropic_cache_control_hook_multiple_content_items_last_only():
|
||||
client = AsyncHTTPHandler()
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
@@ -874,7 +874,7 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages():
|
||||
client = AsyncHTTPHandler()
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
@@ -1057,7 +1057,7 @@ async def test_anthropic_cache_control_hook_string_negative_index():
|
||||
client = AsyncHTTPHandler()
|
||||
with patch.object(client, "post", return_value=mock_response) as mock_post:
|
||||
await litellm.acompletion(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
messages=[
|
||||
{"role": "user", "content": "First message"},
|
||||
{"role": "assistant", "content": "First response"},
|
||||
|
||||
@@ -262,7 +262,7 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase):
|
||||
class TestOpenTelemetry(unittest.TestCase):
|
||||
POLL_INTERVAL = 0.05
|
||||
POLL_TIMEOUT = 2.0
|
||||
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
HERE = os.path.dirname(__file__)
|
||||
|
||||
@patch.dict(os.environ, {}, clear=True)
|
||||
|
||||
@@ -369,7 +369,7 @@ def test_generic_cost_per_token_gpt55():
|
||||
|
||||
|
||||
def test_generic_cost_per_token_gpt55_pro():
|
||||
"""gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input."""
|
||||
"""gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input."""
|
||||
model = "gpt-5.5-pro"
|
||||
custom_llm_provider = "openai"
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
@@ -378,18 +378,18 @@ def test_generic_cost_per_token_gpt55_pro():
|
||||
model_cost_map = litellm.model_cost[model]
|
||||
|
||||
# Sanity-check the map values match OpenAI's published pricing.
|
||||
assert model_cost_map["input_cost_per_token"] == 6e-5
|
||||
assert model_cost_map["output_cost_per_token"] == 3.6e-4
|
||||
assert model_cost_map["cache_read_input_token_cost"] == 6e-6
|
||||
assert model_cost_map["input_cost_per_token"] == 3e-5
|
||||
assert model_cost_map["output_cost_per_token"] == 1.8e-4
|
||||
assert model_cost_map["cache_read_input_token_cost"] == 3e-6
|
||||
assert model_cost_map["litellm_provider"] == "openai"
|
||||
# gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint).
|
||||
assert model_cost_map["mode"] == "responses"
|
||||
assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"]
|
||||
assert "/v1/responses" in model_cost_map["supported_endpoints"]
|
||||
# Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x).
|
||||
# Inherits GPT-5.4-pro's long-context window + tiered pricing.
|
||||
assert model_cost_map["max_input_tokens"] == 1050000
|
||||
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4
|
||||
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4
|
||||
assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5
|
||||
assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4
|
||||
|
||||
prompt_tokens = 1000
|
||||
completion_tokens = 500
|
||||
@@ -454,8 +454,8 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
|
||||
[
|
||||
("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7),
|
||||
("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7),
|
||||
("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6),
|
||||
("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6),
|
||||
("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6),
|
||||
("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6),
|
||||
],
|
||||
)
|
||||
def test_azure_gpt55_entries_present_with_correct_pricing(
|
||||
@@ -464,7 +464,7 @@ def test_azure_gpt55_entries_present_with_correct_pricing(
|
||||
"""Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure.
|
||||
|
||||
Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page
|
||||
on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro.
|
||||
on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro.
|
||||
Cache discount is 10% of input.
|
||||
"""
|
||||
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
|
||||
|
||||
+1
-1
@@ -77,7 +77,7 @@ async def test_anthropic_bedrock_thinking_blocks_with_none_content():
|
||||
# test _bedrock_converse_messages_pt_async
|
||||
result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
|
||||
messages=messages,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
"""
|
||||
Regression tests for the parsed-URL hostname match used to identify a
|
||||
caller-supplied ``api_base`` as a known openai-compatible provider.
|
||||
|
||||
The previous shape (``if endpoint in api_base:``) used unanchored
|
||||
substring search, which let a caller pass
|
||||
``https://attacker.com/api.groq.com/openai/v1`` and have the proxy
|
||||
return ``GROQ_API_KEY`` as the dynamic credential — exfiltrating the
|
||||
server's real provider key to an attacker-controlled host on the
|
||||
outbound request.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
_endpoint_matches_api_base,
|
||||
get_llm_provider,
|
||||
)
|
||||
|
||||
|
||||
class TestEndpointMatchesApiBase:
|
||||
"""Direct unit tests on the parsed-URL matcher."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint, api_base",
|
||||
[
|
||||
# Bare hostname endpoint, exact host match.
|
||||
("api.perplexity.ai", "https://api.perplexity.ai/v1"),
|
||||
# Endpoint includes a path; api_base path starts with it.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"),
|
||||
# Endpoint with full URL scheme.
|
||||
("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"),
|
||||
# Trailing-slash on registered endpoint must not break match.
|
||||
("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"),
|
||||
# Case-insensitive on hostname.
|
||||
("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"),
|
||||
],
|
||||
)
|
||||
def test_legitimate_provider_urls_match(self, endpoint, api_base):
|
||||
assert _endpoint_matches_api_base(endpoint, api_base) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint, api_base",
|
||||
[
|
||||
# Attacker host, registered endpoint smuggled into path.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://attacker.com/api.groq.com/openai/v1",
|
||||
),
|
||||
# Attacker host, registered endpoint smuggled into a path segment.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://attacker.com/foo/api.groq.com/openai/v1",
|
||||
),
|
||||
# Lookalike host that contains the registered host as a suffix label.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com.attacker.com/openai/v1",
|
||||
),
|
||||
# Lookalike host with the registered host as a prefix.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com.evil.example/openai/v1",
|
||||
),
|
||||
# Right host, wrong path — endpoint requires ``/openai/v1`` prefix.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/v1"),
|
||||
# Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"),
|
||||
# Userinfo / @-injection trick — the ``hostname`` after ``@`` is
|
||||
# what httpx connects to.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com@attacker.com/openai/v1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_attacker_smuggling_does_not_match(self, endpoint, api_base):
|
||||
assert _endpoint_matches_api_base(endpoint, api_base) is False
|
||||
|
||||
|
||||
class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
|
||||
"""
|
||||
End-to-end: ``get_llm_provider`` must NOT return the server's stored
|
||||
secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is
|
||||
attacker-controlled, even when the registered endpoint string appears
|
||||
elsewhere in the URL.
|
||||
"""
|
||||
|
||||
def test_attacker_host_does_not_yield_groq_secret(self):
|
||||
# The function may either fall through (different provider) or
|
||||
# raise BadRequestError because the model can't be identified.
|
||||
# The invariant under test is that ``GROQ_API_KEY`` is never
|
||||
# looked up against an attacker-controlled hostname.
|
||||
import litellm
|
||||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
|
||||
return_value="server-real-groq-key",
|
||||
) as mocked_secret:
|
||||
try:
|
||||
_, _, dynamic_api_key, _ = get_llm_provider(
|
||||
model="some-model",
|
||||
api_base="https://attacker.com/api.groq.com/openai/v1",
|
||||
)
|
||||
# If it returned, the dynamic key must not be the secret.
|
||||
assert dynamic_api_key != "server-real-groq-key"
|
||||
except litellm.exceptions.BadRequestError:
|
||||
# Acceptable outcome: provider unidentifiable, no secret
|
||||
# was returned.
|
||||
pass
|
||||
|
||||
# Regardless of return / raise, the secret must never have been
|
||||
# read against this attacker-controlled api_base.
|
||||
groq_lookups = [
|
||||
call
|
||||
for call in mocked_secret.call_args_list
|
||||
if call.args and call.args[0] == "GROQ_API_KEY"
|
||||
]
|
||||
assert groq_lookups == []
|
||||
|
||||
def test_legitimate_groq_api_base_still_resolves(self):
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
|
||||
return_value="server-real-groq-key",
|
||||
):
|
||||
_, provider, dynamic_api_key, _ = get_llm_provider(
|
||||
model="some-model",
|
||||
api_base="https://api.groq.com/openai/v1",
|
||||
)
|
||||
|
||||
assert provider == "groq"
|
||||
assert dynamic_api_key == "server-real-groq-key"
|
||||
@@ -2337,6 +2337,104 @@ def test_merge_hidden_params_from_response_into_metadata_populates_metadata():
|
||||
assert meta["hidden_params"]["model_id"] == "mid-test"
|
||||
|
||||
|
||||
def test_merge_hidden_params_from_response_into_metadata_backfills_response_cost():
|
||||
"""Streaming metadata should include the already-calculated response cost."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="merge-hp-cost-test",
|
||||
function_id="merge-hp-cost-fn",
|
||||
)
|
||||
logging_obj.model_call_details = {
|
||||
"litellm_params": {"metadata": {}},
|
||||
"response_cost": 0.002,
|
||||
}
|
||||
|
||||
class _Resp:
|
||||
_hidden_params = {"response_cost": None, "model_id": "mid-test"}
|
||||
|
||||
response = _Resp()
|
||||
logging_obj._merge_hidden_params_from_response_into_metadata(response)
|
||||
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
assert meta["hidden_params"]["response_cost"] == 0.002
|
||||
assert meta["hidden_params"]["model_id"] == "mid-test"
|
||||
assert response._hidden_params["response_cost"] is None
|
||||
|
||||
|
||||
def test_standard_logging_hidden_params_backfills_response_cost_without_mutating_response():
|
||||
"""Streaming standard logging payload should expose the calculated response cost."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="standard-hp-cost-test",
|
||||
function_id="standard-hp-cost-fn",
|
||||
)
|
||||
logging_obj.model_call_details = {
|
||||
"litellm_params": {"metadata": {}, "proxy_server_request": {}},
|
||||
"litellm_call_id": "standard-hp-cost-test",
|
||||
"call_type": "acompletion",
|
||||
"stream": True,
|
||||
"model": "gpt-4o-mini",
|
||||
"custom_llm_provider": "openai",
|
||||
"optional_params": {"stream": True},
|
||||
"response_cost": 0.002,
|
||||
}
|
||||
response = ModelResponse(
|
||||
id="standard-hp-cost-response",
|
||||
model="gpt-4o-mini",
|
||||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
)
|
||||
response._hidden_params = {"response_cost": None, "model_id": "mid-test"}
|
||||
|
||||
payload = logging_obj._build_standard_logging_payload(
|
||||
response, datetime.now(), datetime.now()
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["hidden_params"]["response_cost"] == 0.002
|
||||
assert response._hidden_params["response_cost"] is None
|
||||
|
||||
|
||||
def test_merge_hidden_params_from_response_into_metadata_preserves_response_cost():
|
||||
"""Do not overwrite provider-supplied response cost when it already exists."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="acompletion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="merge-hp-preserve-cost-test",
|
||||
function_id="merge-hp-preserve-cost-fn",
|
||||
)
|
||||
logging_obj.model_call_details = {
|
||||
"litellm_params": {"metadata": {}},
|
||||
"response_cost": 0.002,
|
||||
}
|
||||
|
||||
class _Resp:
|
||||
_hidden_params = {"response_cost": 0.001, "model_id": "mid-test"}
|
||||
|
||||
logging_obj._merge_hidden_params_from_response_into_metadata(_Resp())
|
||||
meta = logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
assert meta["hidden_params"]["response_cost"] == 0.001
|
||||
assert meta["hidden_params"]["model_id"] == "mid-test"
|
||||
|
||||
|
||||
def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty():
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
@@ -2436,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
|
||||
|
||||
assert payload is not None
|
||||
assert payload["litellm_call_id"] == call_id
|
||||
|
||||
|
||||
def _make_dict_logging_obj():
|
||||
"""Build a Logging instance configured for a non-streaming dict result."""
|
||||
obj = LitellmLogging(
|
||||
model="claude-haiku-4-5@20251001",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
litellm_call_id="test-call-id",
|
||||
start_time=time.time(),
|
||||
function_id="test-fn",
|
||||
)
|
||||
obj.model_call_details = {
|
||||
"model": "claude-haiku-4-5@20251001",
|
||||
"custom_llm_provider": "vertex_ai",
|
||||
"litellm_params": {"metadata": {}},
|
||||
"response_cost": None,
|
||||
}
|
||||
return obj
|
||||
|
||||
|
||||
def test_success_handler_computes_cost_for_dict_response():
|
||||
"""Non-streaming dict responses run through the cost calculator."""
|
||||
logging_obj = _make_dict_logging_obj()
|
||||
expected_cost = 0.42
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_response_cost_calculator",
|
||||
return_value=expected_cost,
|
||||
) as mock_calc,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": expected_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_transform_usage_objects",
|
||||
side_effect=lambda result: result,
|
||||
),
|
||||
):
|
||||
logging_obj.success_handler(
|
||||
result={"id": "msg_1"},
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
mock_calc.assert_called_once()
|
||||
assert logging_obj.model_call_details["response_cost"] == expected_cost
|
||||
|
||||
|
||||
def test_success_handler_preserves_precomputed_cost_for_dict_response():
|
||||
"""Precomputed response_cost on model_call_details must not be overwritten."""
|
||||
logging_obj = _make_dict_logging_obj()
|
||||
precomputed_cost = 1.23
|
||||
logging_obj.model_call_details["response_cost"] = precomputed_cost
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_response_cost_calculator",
|
||||
return_value=9.99,
|
||||
) as mock_calc,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": precomputed_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
return_value=False,
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_transform_usage_objects",
|
||||
side_effect=lambda result: result,
|
||||
),
|
||||
):
|
||||
logging_obj.success_handler(
|
||||
result={"id": "msg_2"},
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
mock_calc.assert_not_called()
|
||||
assert logging_obj.model_call_details["response_cost"] == precomputed_cost
|
||||
|
||||
|
||||
def test_success_handler_unified_helper_runs_for_typed_results():
|
||||
"""Recognized typed responses still flow through the unified helper."""
|
||||
logging_obj = _make_dict_logging_obj()
|
||||
expected_cost = 0.10
|
||||
typed_result = MagicMock()
|
||||
typed_result._hidden_params = {}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_response_cost_calculator",
|
||||
return_value=expected_cost,
|
||||
) as mock_calc,
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_build_standard_logging_payload",
|
||||
return_value={"response_cost": expected_cost},
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_is_recognized_call_type_for_logging",
|
||||
return_value=True,
|
||||
),
|
||||
patch.object(
|
||||
logging_obj,
|
||||
"_transform_usage_objects",
|
||||
side_effect=lambda result: result,
|
||||
),
|
||||
):
|
||||
logging_obj.success_handler(
|
||||
result=typed_result,
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
mock_calc.assert_called_once()
|
||||
assert logging_obj.model_call_details["response_cost"] == expected_cost
|
||||
|
||||
@@ -279,7 +279,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto():
|
||||
}
|
||||
|
||||
optional_params = config.map_openai_params(
|
||||
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
drop_params=False,
|
||||
@@ -2797,7 +2797,7 @@ def test_thinking_with_max_completion_tokens():
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_with_max_completion,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
@@ -2819,7 +2819,7 @@ def test_thinking_with_max_completion_tokens():
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_with_max_tokens,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
@@ -2842,7 +2842,7 @@ def test_thinking_with_max_completion_tokens():
|
||||
result = config.map_openai_params(
|
||||
non_default_params=non_default_params_without_max,
|
||||
optional_params=optional_params,
|
||||
model="us.anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
@@ -3617,7 +3617,7 @@ class TestBedrockMinThinkingBudgetTokens:
|
||||
"""Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024)."""
|
||||
|
||||
def _map_params(
|
||||
self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0"
|
||||
self, thinking_value, model="anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
):
|
||||
"""Helper to call map_openai_params with the given thinking value."""
|
||||
config = AmazonConverseConfig()
|
||||
@@ -3651,7 +3651,7 @@ class TestBedrockMinThinkingBudgetTokens:
|
||||
result = config.map_openai_params(
|
||||
non_default_params={},
|
||||
optional_params={},
|
||||
model="anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
model="anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
drop_params=False,
|
||||
)
|
||||
assert "thinking" not in result or result.get("thinking") is None
|
||||
@@ -4146,3 +4146,39 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools(
|
||||
|
||||
# finish_reason must be "stop", not "tool_calls"
|
||||
assert result.choices[0].finish_reason == "stop"
|
||||
|
||||
|
||||
def test_transform_response_does_not_leak_body_on_parse_failure():
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
||||
leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}}
|
||||
|
||||
class MockResponse:
|
||||
def json(self):
|
||||
return leaky_body
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
return json.dumps(leaky_body)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock",
|
||||
side_effect=KeyError("missing required field"),
|
||||
):
|
||||
with pytest.raises(BedrockError) as exc_info:
|
||||
AmazonConverseConfig()._transform_response(
|
||||
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
response=MockResponse(),
|
||||
model_response=ModelResponse(),
|
||||
stream=False,
|
||||
logging_obj=None,
|
||||
optional_params={},
|
||||
api_key=None,
|
||||
data=None,
|
||||
messages=[],
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
msg = str(exc_info.value)
|
||||
assert "secret content" not in msg
|
||||
assert "Error converting to valid response block" in msg
|
||||
|
||||
+105
@@ -21,7 +21,112 @@ def test_transform_search_request():
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
assert url.endswith("/kb123/retrieve")
|
||||
assert body["retrievalQuery"].get("text") == "hello"
|
||||
|
||||
|
||||
def test_transform_search_request_uses_only_retrieval_config_from_extra_body():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
|
||||
url, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body={
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"overrideSearchType": "HYBRID",
|
||||
"numberOfResults": 8,
|
||||
}
|
||||
},
|
||||
"unrelatedField": {"should_not": "be_forwarded"},
|
||||
},
|
||||
)
|
||||
|
||||
assert url.endswith("/kb123/retrieve")
|
||||
assert body["retrievalQuery"].get("text") == "hello"
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"][
|
||||
"overrideSearchType"
|
||||
]
|
||||
== "HYBRID"
|
||||
)
|
||||
assert "unrelatedField" not in body
|
||||
|
||||
|
||||
def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
extra_body = {
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"overrideSearchType": "HYBRID",
|
||||
"numberOfResults": 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={"max_num_results": 10},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"]
|
||||
== 10
|
||||
)
|
||||
assert (
|
||||
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][
|
||||
"numberOfResults"
|
||||
]
|
||||
== 8
|
||||
)
|
||||
|
||||
|
||||
def test_transform_search_request_overrides_filter_without_mutating_extra_body():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
extra_body = {
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"filter": {"equals": {"key": "tenant", "value": "a"}}
|
||||
}
|
||||
}
|
||||
}
|
||||
new_filter = {"equals": {"key": "tenant", "value": "b"}}
|
||||
|
||||
_, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={"filters": new_filter},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"]
|
||||
== new_filter
|
||||
)
|
||||
assert (
|
||||
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][
|
||||
"equals"
|
||||
]["value"]
|
||||
== "a"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
import socket
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def _invoke_connector_factory(http_handler_module):
|
||||
"""
|
||||
Drive the lambda factory installed on the transport so TCPConnector is
|
||||
actually constructed. _create_aiohttp_transport returns a transport whose
|
||||
_client_factory is the lambda that builds (TCPConnector → ClientSession);
|
||||
invoking it directly avoids relying on _get_valid_client_session's internal
|
||||
branching to trigger connector construction.
|
||||
"""
|
||||
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
|
||||
shared_session=None
|
||||
)
|
||||
transport._client_factory()
|
||||
return transport
|
||||
|
||||
|
||||
def test_socket_factory_omitted_when_disabled(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
|
||||
|
||||
|
||||
def test_socket_factory_attached_when_enabled(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
factory = mock_tcp_connector.call_args.kwargs.get("socket_factory")
|
||||
assert callable(factory)
|
||||
|
||||
|
||||
def test_socket_factory_skipped_on_old_aiohttp(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
|
||||
|
||||
|
||||
def test_socket_factory_sets_keepalive_options(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4)
|
||||
|
||||
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
|
||||
assert factory is not None
|
||||
|
||||
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
|
||||
|
||||
fake_sock = MagicMock(spec=socket.socket)
|
||||
with patch("socket.socket", return_value=fake_sock) as sock_ctor:
|
||||
returned = factory(addr_info)
|
||||
|
||||
sock_ctor.assert_called_once_with(
|
||||
family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
|
||||
)
|
||||
assert returned is fake_sock
|
||||
fake_sock.setblocking.assert_called_once_with(False)
|
||||
|
||||
setsockopt_calls = {
|
||||
(call.args[0], call.args[1]): call.args[2]
|
||||
for call in fake_sock.setsockopt.call_args_list
|
||||
}
|
||||
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
|
||||
|
||||
if hasattr(socket, "TCP_KEEPIDLE"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45
|
||||
elif hasattr(socket, "TCP_KEEPALIVE"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45
|
||||
if hasattr(socket, "TCP_KEEPINTVL"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15
|
||||
if hasattr(socket, "TCP_KEEPCNT"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4
|
||||
|
||||
|
||||
def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch):
|
||||
"""
|
||||
Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE
|
||||
is present, the factory should fall back to TCP_KEEPALIVE for the idle timer.
|
||||
Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to
|
||||
simulate the BSD-derived environment.
|
||||
"""
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60)
|
||||
|
||||
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
|
||||
assert factory is not None
|
||||
|
||||
fake_socket_module = MagicMock(spec=[])
|
||||
fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET
|
||||
fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE
|
||||
fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP
|
||||
fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
|
||||
fake_sock = MagicMock(spec=socket.socket)
|
||||
fake_socket_module.socket = MagicMock(return_value=fake_sock)
|
||||
|
||||
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
|
||||
|
||||
with patch.object(http_handler_module, "socket", fake_socket_module):
|
||||
factory(addr_info)
|
||||
|
||||
setsockopt_calls = {
|
||||
(call.args[0], call.args[1]): call.args[2]
|
||||
for call in fake_sock.setsockopt.call_args_list
|
||||
}
|
||||
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
|
||||
assert (
|
||||
setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
|
||||
)
|
||||
assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls
|
||||
@@ -2,12 +2,16 @@ import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
BaseLLMHTTPHandler,
|
||||
_google_genai_streaming_hidden_params,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
||||
@@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority():
|
||||
assert captured_headers["X-Forwarded-Only"] == "keep"
|
||||
assert captured_headers["X-Extra-Only"] == "also-keep"
|
||||
assert captured_headers["X-Provider-Only"] == "keep-this-too"
|
||||
|
||||
|
||||
def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
|
||||
logging_obj = Mock()
|
||||
logging_obj.get_router_model_id = Mock(return_value="router-model-id")
|
||||
|
||||
from_model_info = _google_genai_streaming_hidden_params(
|
||||
api_base="https://generativelanguage.googleapis.com/v1beta",
|
||||
litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}),
|
||||
logging_obj=logging_obj,
|
||||
response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}),
|
||||
)
|
||||
assert from_model_info["model_id"] == "info-id"
|
||||
assert (
|
||||
from_model_info["api_base"]
|
||||
== "https://generativelanguage.googleapis.com/v1beta"
|
||||
)
|
||||
assert isinstance(from_model_info["additional_headers"], dict)
|
||||
|
||||
from_router = _google_genai_streaming_hidden_params(
|
||||
api_base="https://x",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
logging_obj=logging_obj,
|
||||
response_headers=httpx.Headers({}),
|
||||
)
|
||||
assert from_router["model_id"] == "router-model-id"
|
||||
|
||||
@@ -59,6 +59,7 @@ class TestS3VectorsVectorStoreConfig:
|
||||
api_base="https://s3vectors.us-west-2.api.aws",
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
def test_transform_search_response(self):
|
||||
|
||||
@@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client():
|
||||
# Verify that gemini_client is in the partial's keywords
|
||||
assert "gemini_client" in partial_make_sync_call.keywords
|
||||
assert partial_make_sync_call.keywords["gemini_client"] is mock_client
|
||||
|
||||
|
||||
def test_transform_response_does_not_leak_body_on_parse_failure():
|
||||
leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]}
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = leaky_body
|
||||
raw_response.text = json.dumps(leaky_body)
|
||||
raw_response.headers = {}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody",
|
||||
side_effect=KeyError("missing required field"),
|
||||
):
|
||||
with pytest.raises(VertexAIError) as exc_info:
|
||||
VertexGeminiConfig().transform_response(
|
||||
model="gemini-pro",
|
||||
raw_response=raw_response,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
msg = str(exc_info.value)
|
||||
assert "secret content" not in msg
|
||||
assert "Error converting to valid response block" in msg
|
||||
|
||||
@@ -225,7 +225,11 @@ def test_build_vertex_schema():
|
||||
"metadata": {"type": "object"},
|
||||
"callbacks": {
|
||||
"anyOf": [
|
||||
{"type": "array", "nullable": True},
|
||||
{
|
||||
"type": "array",
|
||||
"items": {"type": "object"},
|
||||
"nullable": True,
|
||||
},
|
||||
{"type": "object", "nullable": True},
|
||||
]
|
||||
},
|
||||
@@ -288,6 +292,43 @@ def test_process_items_basic():
|
||||
process_items(schema)
|
||||
assert schema["properties"]["nested"]["items"] == {"type": "object"}
|
||||
|
||||
# Vertex rejects array types missing `items` entirely (not just empty).
|
||||
# Synthesize {"type": "object"} so the request validates.
|
||||
schema = {"type": "array"}
|
||||
process_items(schema)
|
||||
assert schema["items"] == {"type": "object"}
|
||||
|
||||
|
||||
def test_build_vertex_schema_array_branch_missing_items_in_anyof():
|
||||
"""
|
||||
Regression: an `anyOf` branch with `{"type": "array"}` (no items) must
|
||||
end up with synthesized `items: {"type": "object"}` after the schema
|
||||
transform — Vertex returns INVALID_ARGUMENT otherwise.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
|
||||
|
||||
parameters = {
|
||||
"properties": {
|
||||
"callbacks": {
|
||||
"anyOf": [
|
||||
{"type": "array"},
|
||||
{"type": "object"},
|
||||
{"type": "null"},
|
||||
]
|
||||
}
|
||||
},
|
||||
"type": "object",
|
||||
}
|
||||
|
||||
result = _build_vertex_schema(parameters)
|
||||
callbacks_anyof = result["properties"]["callbacks"]["anyOf"]
|
||||
array_branches = [b for b in callbacks_anyof if b.get("type") == "array"]
|
||||
assert array_branches, "expected an array branch to remain after transform"
|
||||
for branch in array_branches:
|
||||
assert branch.get("items") == {
|
||||
"type": "object"
|
||||
}, f"array branch must have items synthesized; got {branch}"
|
||||
|
||||
|
||||
def test_vertex_ai_complex_response_schema():
|
||||
import json
|
||||
|
||||
+26
@@ -311,3 +311,29 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control()
|
||||
# scope removed from message content
|
||||
assert "scope" not in result["messages"][0]["content"][0]["cache_control"]
|
||||
assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral"
|
||||
|
||||
|
||||
def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance():
|
||||
"""
|
||||
Regression test: repeated provider config lookups for the same Vertex Claude model
|
||||
should return the same config instance (which preserves auth cache state).
|
||||
"""
|
||||
import litellm
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear()
|
||||
try:
|
||||
first_config = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model="claude-opus-4-6",
|
||||
provider=litellm.LlmProviders.VERTEX_AI,
|
||||
)
|
||||
second_config = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model="claude-opus-4-6",
|
||||
provider=litellm.LlmProviders.VERTEX_AI,
|
||||
)
|
||||
|
||||
assert isinstance(first_config, VertexAIPartnerModelsAnthropicMessagesConfig)
|
||||
assert isinstance(second_config, VertexAIPartnerModelsAnthropicMessagesConfig)
|
||||
assert first_config is second_config
|
||||
finally:
|
||||
ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear()
|
||||
|
||||
@@ -728,6 +728,136 @@ class TestMCPServerManager:
|
||||
]
|
||||
assert scopes == ["read", "write"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_metadata = MCPOAuthMetadata(
|
||||
scopes=None,
|
||||
authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize",
|
||||
token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token",
|
||||
registration_url=None,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
),
|
||||
patch.object(
|
||||
manager,
|
||||
"_attempt_well_known_discovery",
|
||||
AsyncMock(
|
||||
return_value=(
|
||||
["https://login.microsoftonline.com/test-tenant-id/v2.0"],
|
||||
["api://some-scope/.default"],
|
||||
)
|
||||
),
|
||||
) as mock_well_known,
|
||||
patch.object(
|
||||
manager,
|
||||
"_fetch_authorization_server_metadata",
|
||||
AsyncMock(return_value=mock_metadata),
|
||||
) as mock_fetch_auth,
|
||||
):
|
||||
result = await manager._descovery_metadata("http://localhost:8001/mcp")
|
||||
|
||||
mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp")
|
||||
mock_fetch_auth.assert_awaited_once_with(
|
||||
["https://login.microsoftonline.com/test-tenant-id/v2.0"]
|
||||
)
|
||||
assert result is mock_metadata
|
||||
assert result.scopes == ["api://some-scope/.default"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
|
||||
|
||||
def build_response(url: str):
|
||||
mock_response = MagicMock()
|
||||
if url == f"{issuer}/.well-known/openid-configuration":
|
||||
mock_response.json.return_value = {
|
||||
"authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize",
|
||||
"token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token",
|
||||
"scopes_supported": ["api://some-scope/.default"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
else:
|
||||
request = httpx.Request("GET", url)
|
||||
response_obj = httpx.Response(status_code=404, request=request)
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"not found", request=request, response=response_obj
|
||||
)
|
||||
)
|
||||
return mock_response
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(side_effect=build_response)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await manager._fetch_single_authorization_server_metadata(issuer)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result.authorization_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
|
||||
)
|
||||
assert (
|
||||
result.token_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
|
||||
)
|
||||
assert result.scopes == ["api://some-scope/.default"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_single_authorization_server_metadata_derives_azure_metadata(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
|
||||
|
||||
request = httpx.Request("GET", issuer)
|
||||
response_obj = httpx.Response(status_code=404, request=request)
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"not found", request=request, response=response_obj
|
||||
)
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await manager._fetch_single_authorization_server_metadata(issuer)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result.authorization_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
|
||||
)
|
||||
assert (
|
||||
result.token_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self):
|
||||
manager = MCPServerManager()
|
||||
|
||||
@@ -0,0 +1,302 @@
|
||||
"""
|
||||
Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX).
|
||||
|
||||
The short-prefix mode swaps the historical alias/server_name prefix on
|
||||
tool names for a deterministic three-character base62 ID derived from the
|
||||
server's ``server_id``. This keeps tool names well below the 60-char
|
||||
upper bound enforced by some model APIs while remaining stable across
|
||||
processes/restarts and tolerant of mixed-version clients.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH,
|
||||
add_server_prefix_to_name,
|
||||
compute_short_server_prefix,
|
||||
get_server_prefix,
|
||||
is_short_mcp_tool_prefix_enabled,
|
||||
iter_known_server_prefixes,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
def _make_server(
|
||||
*,
|
||||
server_id: str = "abcdef-1234",
|
||||
server_name: str = "github_onprem",
|
||||
alias: str = "github_onprem",
|
||||
) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=alias or server_name,
|
||||
alias=alias,
|
||||
server_name=server_name,
|
||||
transport="http",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_env(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False)
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixHelpers:
|
||||
def test_short_prefix_is_three_base62_chars(self):
|
||||
prefix = compute_short_server_prefix("any-server-id")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
assert prefix.isalnum() and prefix.isascii()
|
||||
|
||||
def test_short_prefix_first_char_is_alphabetic(self):
|
||||
"""The first char must be [A-Za-z] so the prefix is a valid identifier
|
||||
on every model API (some providers historically required the first
|
||||
character of a function name to be alphabetic)."""
|
||||
# Sweep many server_ids and rehash attempts to give us coverage of
|
||||
# every position the high-order bits can land on.
|
||||
for i in range(200):
|
||||
for attempt in range(4):
|
||||
prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt)
|
||||
assert prefix[0].isalpha(), (
|
||||
f"prefix {prefix!r} for server-{i} (attempt={attempt}) "
|
||||
f"starts with a non-alphabetic character"
|
||||
)
|
||||
|
||||
def test_short_prefix_is_deterministic(self):
|
||||
assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc")
|
||||
assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd")
|
||||
|
||||
def test_short_prefix_requires_server_id(self):
|
||||
with pytest.raises(ValueError):
|
||||
compute_short_server_prefix("")
|
||||
|
||||
def test_flag_defaults_to_false(self):
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"])
|
||||
def test_flag_truthy_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is True
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
|
||||
def test_flag_falsey_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_server_prefix behaviour
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetServerPrefix:
|
||||
def test_default_mode_uses_alias(self):
|
||||
server = _make_server(alias="github_onprem", server_name="github_onprem")
|
||||
assert get_server_prefix(server) == "github_onprem"
|
||||
|
||||
def test_short_mode_uses_short_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abcdef-1234")
|
||||
prefix = get_server_prefix(server)
|
||||
assert prefix == compute_short_server_prefix("abcdef-1234")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
|
||||
def test_short_mode_falls_back_when_no_server_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
class _Bare:
|
||||
alias = "fallback_alias"
|
||||
server_name = None
|
||||
server_id = None
|
||||
|
||||
assert get_server_prefix(_Bare()) == "fallback_alias"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# iter_known_server_prefixes — covers reverse-lookup tolerance
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIterKnownServerPrefixes:
|
||||
def test_default_mode_includes_short_id_too(self):
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
# Contains the live prefix and every known form so that mixed-mode
|
||||
# clients can be resolved.
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
def test_short_mode_still_yields_long_forms(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Manager-level behaviour: list + reverse-lookup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_tools() -> List[MCPTool]:
|
||||
return [
|
||||
MCPTool(name="get_repo", description="", inputSchema={"type": "object"}),
|
||||
MCPTool(name="list_issues", description="", inputSchema={"type": "object"}),
|
||||
]
|
||||
|
||||
|
||||
class TestManagerShortPrefix:
|
||||
def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"}
|
||||
|
||||
def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch):
|
||||
"""Old clients that cached the long-prefix name must still route."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_default_mode_unchanged(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
assert {t.name for t in out} == {
|
||||
"github_onprem-get_repo",
|
||||
"github_onprem-list_issues",
|
||||
}
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None
|
||||
) # registry empty
|
||||
manager.registry[server.server_id] = server
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server
|
||||
)
|
||||
|
||||
def test_total_tool_name_length_short_enough(self, monkeypatch):
|
||||
"""The short prefix keeps tool names under the 60-char limit even
|
||||
when the upstream tool name is itself reasonably long."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
long_server_name = "a" * 50
|
||||
server = _make_server(
|
||||
server_id="server-id-1",
|
||||
server_name=long_server_name,
|
||||
alias=long_server_name,
|
||||
)
|
||||
prefix = get_server_prefix(server)
|
||||
full = add_server_prefix_to_name("get_repo", prefix)
|
||||
assert len(full) < 60
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collision-resolution at registration time
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixCollisionResolution:
|
||||
"""``_assign_unique_short_prefix`` must rehash on collision.
|
||||
|
||||
The dedup path is exercised by forcing two distinct ``server_id``
|
||||
values to both hash to the same natural prefix via a monkeypatched
|
||||
``compute_short_server_prefix``.
|
||||
"""
|
||||
|
||||
def test_no_op_when_flag_off(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
assert server.short_prefix is None
|
||||
|
||||
def test_assigns_natural_hash_when_no_collision(self, monkeypatch):
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc")
|
||||
|
||||
def test_rehashes_when_natural_hash_collides(self, monkeypatch):
|
||||
"""Two server_ids that natural-hash to the same prefix get
|
||||
deterministic, distinct short prefixes."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
# Force every attempt=0 hash to "AAA" and attempt=1 to "AAB".
|
||||
# That way the second server registered must rehash to "AAB".
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
def _fake_hash(server_id: str, attempt: int = 0) -> str:
|
||||
return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}"
|
||||
|
||||
monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash)
|
||||
# Also patch the symbol that the manager imported at module load.
|
||||
from litellm.proxy._experimental.mcp_server import (
|
||||
mcp_server_manager as mgr_module,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash)
|
||||
|
||||
manager = MCPServerManager()
|
||||
first = _make_server(server_id="server-1", alias="srv1")
|
||||
second = _make_server(server_id="server-2", alias="srv2")
|
||||
|
||||
# Pretend both are already in the registry so dedup sees both.
|
||||
manager.registry[first.server_id] = first
|
||||
manager._assign_unique_short_prefix(first)
|
||||
manager.registry[second.server_id] = second
|
||||
manager._assign_unique_short_prefix(second)
|
||||
|
||||
assert first.short_prefix == "AAA"
|
||||
assert second.short_prefix == "AAB"
|
||||
assert first.short_prefix != second.short_prefix
|
||||
|
||||
def test_cached_prefix_is_reused(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "ZZZ" # pretend a previous registration set this
|
||||
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == "ZZZ"
|
||||
|
||||
def test_get_server_prefix_prefers_cached(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "Q9q"
|
||||
|
||||
assert get_server_prefix(server) == "Q9q"
|
||||
@@ -5,6 +5,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_get_customer_id_from_standard_headers,
|
||||
@@ -15,6 +17,7 @@ from litellm.proxy.auth.auth_utils import (
|
||||
get_key_model_tpm_limit,
|
||||
get_project_model_rpm_limit,
|
||||
get_project_model_tpm_limit,
|
||||
is_request_body_safe,
|
||||
)
|
||||
|
||||
|
||||
@@ -660,3 +663,304 @@ class TestCheckCompleteCredentials:
|
||||
def test_returns_true_when_api_key_is_valid(self):
|
||||
result = check_complete_credentials({"model": "gpt-4", "api_key": "sk-valid"})
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestCheckCompleteCredentialsBlocksSSRF:
|
||||
"""
|
||||
Even with credentials supplied, ``api_base`` / ``base_url`` must not
|
||||
point at private / internal / cloud-metadata addresses. Without this
|
||||
the gate accepts ``api_key=anything`` plus a malicious target and the
|
||||
proxy is used as an SSRF pivot.
|
||||
|
||||
The check only runs when ``litellm.user_url_validation`` is True, so
|
||||
every test in this class flips the toggle. Tests stay mock-only — no
|
||||
real DNS is performed.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_url_validation(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", True, raising=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url_field",
|
||||
["api_base", "base_url"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"blocked_url",
|
||||
[
|
||||
"http://169.254.169.254/latest/meta-data/iam/security-credentials/",
|
||||
"http://metadata.google.internal/computeMetadata/v1/",
|
||||
"http://127.0.0.1:8080/admin",
|
||||
"http://10.0.0.1/",
|
||||
"http://192.168.1.1/",
|
||||
],
|
||||
)
|
||||
def test_rejects_private_or_metadata_targets(self, url_field, blocked_url):
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
side_effect=SSRFError(f"blocked: {blocked_url}"),
|
||||
):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
url_field: blocked_url,
|
||||
}
|
||||
)
|
||||
assert url_field in str(exc_info.value)
|
||||
assert "SSRF" in str(exc_info.value)
|
||||
|
||||
def test_allows_public_target_when_validate_url_passes(self):
|
||||
# ``validate_url`` is mocked so no real DNS is performed.
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
return_value=("https://api.openai.com/v1", "api.openai.com"),
|
||||
):
|
||||
result = check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
}
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_skips_url_validation_when_toggle_is_off(self, monkeypatch):
|
||||
# Admins who disable ``user_url_validation`` (default) should not
|
||||
# have requests rejected at the proxy boundary even if the URL
|
||||
# would fail the SSRF guard.
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
) as mocked:
|
||||
result = check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
"api_base": "http://127.0.0.1:8080/admin",
|
||||
}
|
||||
)
|
||||
assert result is True
|
||||
mocked.assert_not_called()
|
||||
|
||||
|
||||
class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
||||
"""
|
||||
When the caller redirects ``api_base`` / ``base_url`` to their own
|
||||
server, admin-set fields like ``OpenAI-Organization``, ``extra_body``,
|
||||
AWS / Vertex / Azure tokens, and per-deployment ``api_version`` must
|
||||
NOT flow through to that destination.
|
||||
"""
|
||||
|
||||
def test_clears_admin_organization_and_extra_body_on_base_override(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
admin_params = {
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-admin-key",
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin-corp",
|
||||
"extra_body": {"x-admin-secret": "super-secret"},
|
||||
"api_version": "2026-04-01",
|
||||
}
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params=dict(admin_params),
|
||||
request_kwargs={
|
||||
"api_key": "sk-attacker",
|
||||
"api_base": "https://attacker.example",
|
||||
},
|
||||
)
|
||||
assert out["api_base"] == "https://attacker.example"
|
||||
assert out["api_key"] == "sk-attacker"
|
||||
assert "organization" not in out
|
||||
assert "extra_body" not in out
|
||||
assert "api_version" not in out
|
||||
|
||||
def test_clears_aws_and_vertex_secrets_on_base_override(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
admin_params = {
|
||||
"model": "bedrock/claude-3",
|
||||
"aws_access_key_id": "AKIA-EXAMPLE",
|
||||
"aws_secret_access_key": "secret-example",
|
||||
"aws_session_token": "session-example",
|
||||
"vertex_credentials": '{"private_key":"-----BEGIN..."}',
|
||||
"vertex_project": "admin-gcp-project",
|
||||
}
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params=dict(admin_params),
|
||||
request_kwargs={"base_url": "https://attacker.example"},
|
||||
)
|
||||
assert "aws_access_key_id" not in out
|
||||
assert "aws_secret_access_key" not in out
|
||||
assert "aws_session_token" not in out
|
||||
assert "vertex_credentials" not in out
|
||||
assert "vertex_project" not in out
|
||||
|
||||
def test_caller_resupplied_value_overrides_admin_value_on_base_override(self):
|
||||
# When the caller redirects ``api_base`` and *also* supplies their
|
||||
# own value for one of the admin fields (e.g. ``organization``),
|
||||
# the caller's value must win — never the admin's. The naive
|
||||
# ``if field not in request_kwargs: pop`` shape lets a caller echo
|
||||
# the field name with any value (or empty string) to keep the
|
||||
# admin's value forwarded, which is the exfiltration vector this
|
||||
# test guards against.
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin",
|
||||
"extra_body": {"admin": "value"},
|
||||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"organization": "org-attacker",
|
||||
"extra_body": {"attacker": "value"},
|
||||
},
|
||||
)
|
||||
assert out["organization"] == "org-attacker"
|
||||
assert out["extra_body"] == {"attacker": "value"}
|
||||
|
||||
def test_field_echo_does_not_preserve_admin_value(self):
|
||||
# Regression: a caller that echoes an admin-config field name with
|
||||
# an *empty* value (or any value) must not be able to keep the
|
||||
# admin's value in ``litellm_params``.
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin-secret",
|
||||
"extra_body": {"x-admin-only": "secret"},
|
||||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"organization": "",
|
||||
"extra_body": "",
|
||||
},
|
||||
)
|
||||
assert out["organization"] == ""
|
||||
assert out["extra_body"] == ""
|
||||
assert "org-admin-secret" not in str(out)
|
||||
|
||||
def test_no_clearing_when_only_api_key_overridden(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
# Caller only overrides api_key (BYOK pattern); admin's organization /
|
||||
# extra_body / region still apply because the destination is unchanged.
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin",
|
||||
"api_version": "2026-04-01",
|
||||
},
|
||||
request_kwargs={"api_key": "sk-byok"},
|
||||
)
|
||||
assert out["organization"] == "org-admin"
|
||||
assert out["api_version"] == "2026-04-01"
|
||||
assert out["api_base"] == "https://admin.upstream/v1"
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksEndpointTargetingFields:
|
||||
"""
|
||||
``is_request_body_safe`` rejects request-body fields that retarget the
|
||||
outbound request to a caller-controlled host. Beyond the original
|
||||
``api_base`` / ``base_url``, the same protection must apply to:
|
||||
|
||||
* ``aws_bedrock_runtime_endpoint`` — Bedrock endpoint redirect; an
|
||||
attacker-controlled value coerces the proxy to authenticate against
|
||||
their host with the admin's AWS creds.
|
||||
* ``langsmith_base_url`` — Langsmith callback host; attacker-controlled
|
||||
values exfiltrate the entire request payload (incl. message content)
|
||||
via the observability hook.
|
||||
* ``langfuse_host`` — same exfil vector via the Langfuse hook.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _disable_url_validation(self, monkeypatch):
|
||||
# The new banned-params entries should be rejected even when
|
||||
# ``user_url_validation`` is off — the gate isn't the URL guard,
|
||||
# it's the banned-params list.
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
[
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"langsmith_base_url",
|
||||
"langfuse_host",
|
||||
"posthog_host",
|
||||
"braintrust_host",
|
||||
"slack_webhook_url",
|
||||
"s3_endpoint_url",
|
||||
"sagemaker_base_url",
|
||||
"deployment_url",
|
||||
],
|
||||
)
|
||||
def test_endpoint_targeting_field_in_request_body_is_rejected(self, field):
|
||||
with pytest.raises(ValueError) as exc:
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", field: "https://attacker.example"},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
# The function lists the offending param name in the error.
|
||||
assert field in str(exc.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
["api_base", "base_url", "user_config", "langfuse_host", "slack_webhook_url"],
|
||||
)
|
||||
def test_api_key_does_not_bypass_blocklist(self, field):
|
||||
# Regression: the historical ``check_complete_credentials`` clause
|
||||
# made the entire blocklist a no-op for any caller that supplied
|
||||
# a non-empty ``api_key``. That bypass turned every missing entry
|
||||
# on the blocklist into an SSRF / credential-exfil hole. Verify
|
||||
# that supplying an api_key (alongside the banned param) does NOT
|
||||
# bypass the gate — it can only be opened by an admin opt-in.
|
||||
with pytest.raises(ValueError) as exc:
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-anything",
|
||||
field: "https://attacker.example",
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
assert field in str(exc.value)
|
||||
|
||||
def test_admin_opt_in_proxy_wide_still_allows(self):
|
||||
# ``general_settings.allow_client_side_credentials = True`` remains
|
||||
# the documented proxy-wide BYOK opt-in.
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", "api_base": "https://my-byok.example"},
|
||||
general_settings={"allow_client_side_credentials": True},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
@@ -0,0 +1,348 @@
|
||||
"""
|
||||
Test expired UI session key cleanup manager functionality.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.constants import (
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import (
|
||||
ExpiredUISessionKeyCleanupManager,
|
||||
)
|
||||
|
||||
|
||||
class TestExpiredUISessionKeyCleanupManager:
|
||||
"""Test the ExpiredUISessionKeyCleanupManager class functionality."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_find_expired_ui_session_keys_filters_dashboard_team_and_expiry(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
now = datetime(2026, 4, 25, 12, 0, 0, tzinfo=timezone.utc)
|
||||
mock_keys = [
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=now - timedelta(seconds=1),
|
||||
)
|
||||
]
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = (
|
||||
mock_keys
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.datetime"
|
||||
) as mock_datetime:
|
||||
mock_datetime.now.return_value = now
|
||||
mock_datetime.side_effect = lambda *args, **kwargs: datetime(
|
||||
*args, **kwargs
|
||||
)
|
||||
|
||||
keys = await manager._find_expired_ui_session_keys()
|
||||
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with(
|
||||
where={
|
||||
"team_id": UI_SESSION_TOKEN_TEAM_ID,
|
||||
"expires": {"lt": now},
|
||||
},
|
||||
take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
|
||||
)
|
||||
assert keys == mock_keys
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_uses_existing_delete_path(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_key = LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.return_value = (
|
||||
{"deleted_keys": ["expired-dashboard-token"], "failed_tokens": []},
|
||||
[expired_key],
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_key_deleted_hook:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 1
|
||||
mock_delete_verification_tokens.assert_called_once()
|
||||
call_kwargs = mock_delete_verification_tokens.call_args.kwargs
|
||||
assert call_kwargs["tokens"] == ["expired-dashboard-token"]
|
||||
assert call_kwargs["user_api_key_cache"] == mock_cache
|
||||
assert (
|
||||
call_kwargs["litellm_changed_by"]
|
||||
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
|
||||
)
|
||||
assert call_kwargs["user_api_key_dict"].user_id == "system"
|
||||
mock_key_deleted_hook.assert_called_once()
|
||||
hook_kwargs = mock_key_deleted_hook.call_args.kwargs
|
||||
assert hook_kwargs["data"].keys == ["expired-dashboard-token"]
|
||||
assert hook_kwargs["keys_being_deleted"] == [expired_key]
|
||||
assert hook_kwargs["response"] == {
|
||||
"deleted_keys": ["expired-dashboard-token"],
|
||||
"failed_tokens": [],
|
||||
}
|
||||
assert (
|
||||
hook_kwargs["litellm_changed_by"]
|
||||
== LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_deletes_multiple_keys(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_keys = [
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-1",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-2",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
]
|
||||
tokens = [key.token for key in expired_keys]
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.return_value = (
|
||||
{"deleted_keys": tokens, "failed_tokens": []},
|
||||
expired_keys,
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_key_deleted_hook:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 2
|
||||
assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens
|
||||
hook_kwargs = mock_key_deleted_hook.call_args.kwargs
|
||||
assert hook_kwargs["data"].keys == tokens
|
||||
assert hook_kwargs["keys_being_deleted"] == expired_keys
|
||||
assert hook_kwargs["response"] == {"deleted_keys": tokens, "failed_tokens": []}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_returns_successful_delete_count(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_keys = [
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-1",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-2",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
]
|
||||
tokens = [key.token for key in expired_keys]
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.return_value = (
|
||||
{
|
||||
"deleted_keys": ["expired-dashboard-token-1"],
|
||||
"failed_tokens": ["expired-dashboard-token-2"],
|
||||
},
|
||||
[expired_keys[0]],
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
):
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 1
|
||||
assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_counts_nested_delete_response(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_keys = [
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-1",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token-2",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
),
|
||||
]
|
||||
tokens = [key.token for key in expired_keys]
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.return_value = (
|
||||
{
|
||||
"deleted_keys": {"deleted_keys": 2},
|
||||
"failed_tokens": tokens,
|
||||
},
|
||||
expired_keys,
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
):
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_treats_missing_keys_as_noop(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
expired_key = LiteLLM_VerificationToken(
|
||||
token="expired-dashboard-token",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
expires=datetime.now(timezone.utc) - timedelta(seconds=1),
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
mock_delete_verification_tokens.side_effect = HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail={"error": "No keys found"},
|
||||
)
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_key_deleted_hook:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_key_deleted_hook.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_noops_when_no_keys_found(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_delete_verification_tokens:
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_delete_verification_tokens.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_skips_when_lock_held(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
pod_lock_manager=mock_pod_lock_manager,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock()
|
||||
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_pod_lock_manager.acquire_lock.assert_called_once_with(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
manager._find_expired_ui_session_keys.assert_not_called()
|
||||
mock_pod_lock_manager.release_lock.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cleanup_expired_keys_releases_acquired_lock(self):
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_cache = MagicMock()
|
||||
mock_pod_lock_manager = MagicMock()
|
||||
mock_pod_lock_manager.redis_cache = MagicMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
|
||||
manager = ExpiredUISessionKeyCleanupManager(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=mock_cache,
|
||||
pod_lock_manager=mock_pod_lock_manager,
|
||||
)
|
||||
manager._find_expired_ui_session_keys = AsyncMock(return_value=[])
|
||||
|
||||
deleted_count = await manager.cleanup_expired_keys()
|
||||
|
||||
assert deleted_count == 0
|
||||
mock_pod_lock_manager.release_lock.assert_called_once_with(
|
||||
cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
|
||||
)
|
||||
@@ -0,0 +1,255 @@
|
||||
"""
|
||||
Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read,
|
||||
on transport error reconnect once and retry once" helper.
|
||||
|
||||
Covers the regression in issue #25143 where read paths (e.g.
|
||||
`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in
|
||||
LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient
|
||||
`httpx.ReadError` flaps that used to self-heal in 1.82.6.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
|
||||
|
||||
def _make_client(
|
||||
*,
|
||||
attempt_db_reconnect_return: bool = True,
|
||||
has_attempt_db_reconnect: bool = True,
|
||||
):
|
||||
"""Build a minimal stand-in for PrismaClient that exposes only the surface
|
||||
`call_with_db_reconnect_retry` actually pokes at."""
|
||||
client = MagicMock()
|
||||
if has_attempt_db_reconnect:
|
||||
client.attempt_db_reconnect = AsyncMock(
|
||||
return_value=attempt_db_reconnect_return
|
||||
)
|
||||
else:
|
||||
# `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock
|
||||
# auto-creates attributes, so we wipe it out via `spec`.
|
||||
client = MagicMock(spec=[])
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_returns_value_on_first_success():
|
||||
"""Happy path: factory succeeds first call, no reconnect attempted."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path")
|
||||
|
||||
assert result == {"id": 1}
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_retries_after_transport_error():
|
||||
"""Transport error on first call → reconnect → second call succeeds."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="prisma_get_generic_data_config_lookup_failure"
|
||||
)
|
||||
|
||||
assert result == {"id": 1}
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error():
|
||||
"""Data-layer errors (e.g. UniqueViolationError) are NOT transport errors —
|
||||
propagate immediately, do not reconnect."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
raise UniqueViolationError(
|
||||
data={"user_facing_error": {"meta": {}}},
|
||||
message="Unique constraint failed",
|
||||
)
|
||||
|
||||
with pytest.raises(UniqueViolationError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test")
|
||||
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails():
|
||||
"""Transport error, but reconnect returns False → propagate the original
|
||||
exception. Do not call factory a second time."""
|
||||
client = _make_client(attempt_db_reconnect_return=False)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails")
|
||||
|
||||
assert len(invocations) == 1
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error():
|
||||
"""Transport error, reconnect succeeds, retry also raises transport error →
|
||||
propagate. At most one retry by construction (no infinite loop)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("still failing")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="second_transport_error"
|
||||
)
|
||||
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr():
|
||||
"""Older PrismaClient stand-ins / partial mocks may not expose
|
||||
`attempt_db_reconnect`. The helper must not crash — just propagate the
|
||||
original exception. Mirrors the `hasattr` guard from
|
||||
`auth_checks._fetch_key_object_from_db_with_reconnect`."""
|
||||
client = _make_client(has_attempt_db_reconnect=False)
|
||||
|
||||
async def _factory():
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro():
|
||||
"""Guard against the obvious bug of awaiting the same coroutine twice
|
||||
(`RuntimeError: cannot reuse already awaited coroutine`). The helper must
|
||||
call the factory a fresh time on retry, not cache an awaitable."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
factory_call_count = 0
|
||||
|
||||
async def _factory():
|
||||
nonlocal factory_call_count
|
||||
factory_call_count += 1
|
||||
if factory_call_count == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="fresh_coro_on_retry"
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert factory_call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_passes_explicit_timeouts():
|
||||
"""Explicit timeout_seconds / lock_timeout_seconds override the auth
|
||||
defaults read off the prisma_client object."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client,
|
||||
_factory,
|
||||
reason="explicit_timeouts",
|
||||
timeout_seconds=5.5,
|
||||
lock_timeout_seconds=0.25,
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 5.5
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.25
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset():
|
||||
"""When timeouts are not provided, helper reads
|
||||
`_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`
|
||||
off the prisma_client (matching the auth path's existing convention)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
client._db_auth_reconnect_timeout_seconds = 3.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.5
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="defaults")
|
||||
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 3.0
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises():
|
||||
"""If `attempt_db_reconnect` itself raises (lock cancellation, timer
|
||||
error, unexpected internal failure), the helper must surface the
|
||||
*original* transport error to telemetry — not the reconnect exception.
|
||||
Otherwise `failure_handler` / `db_exceptions` alerts log the wrong
|
||||
error string and the actual DB transport problem becomes invisible.
|
||||
|
||||
The reconnect error is chained as the `__cause__` for debuggability."""
|
||||
client = MagicMock()
|
||||
reconnect_exc = RuntimeError("simulated reconnect lock cancellation")
|
||||
client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc)
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
|
||||
original_exc = httpx.ReadError("transport blip")
|
||||
|
||||
async def _factory():
|
||||
raise original_exc
|
||||
|
||||
with pytest.raises(httpx.ReadError) as exc_info:
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="reconnect_itself_raises"
|
||||
)
|
||||
|
||||
assert exc_info.value is original_exc
|
||||
assert exc_info.value.__cause__ is reconnect_exc
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
@@ -5,6 +5,7 @@ import sys
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
@@ -34,18 +35,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_reconnect_success",
|
||||
force=True,
|
||||
)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_reconnect_success",
|
||||
force=True,
|
||||
)
|
||||
|
||||
assert result is True
|
||||
client.db.disconnect.assert_awaited_once()
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
|
||||
|
||||
@@ -140,15 +141,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
|
||||
)
|
||||
client._db_last_reconnect_attempt_ts = 0.0
|
||||
client._db_reconnect_cooldown_seconds = 10
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
# Use a counter-based mock to avoid StopIteration when time.time() is called
|
||||
# more times than expected (varies by Python version / internal code paths).
|
||||
fake_clock = iter(range(100, 10000))
|
||||
with patch(
|
||||
"litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock))
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.utils.time.time",
|
||||
side_effect=lambda: float(next(fake_clock)),
|
||||
),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_cooldown_timestamp_after_attempt",
|
||||
@@ -162,23 +167,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops(
|
||||
async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Direct reconnect goes through recreate_prisma_client (which non-blockingly
|
||||
kills the old engine) instead of calling disconnect() — see issue #26191.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used"))
|
||||
client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used"))
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.disconnect = AsyncMock(
|
||||
side_effect=AssertionError("disconnect must not be called")
|
||||
)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
|
||||
client.db.disconnect.assert_awaited_once()
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -189,19 +199,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._db_watchdog_reconnect_timeout_seconds = 0.1
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
async def _slow_connect():
|
||||
async def _slow_recreate(_db_url):
|
||||
await asyncio.sleep(0.08)
|
||||
|
||||
async def _slow_query(_query: str):
|
||||
await asyncio.sleep(0.08)
|
||||
return [{"result": 1}]
|
||||
|
||||
client.db.connect = AsyncMock(side_effect=_slow_connect)
|
||||
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
|
||||
client.db.query_raw = AsyncMock(side_effect=_slow_query)
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
with (
|
||||
pytest.raises(asyncio.TimeoutError),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
|
||||
|
||||
@@ -212,19 +225,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
async def _slow_connect():
|
||||
async def _slow_recreate(_db_url):
|
||||
await asyncio.sleep(0.08)
|
||||
|
||||
async def _slow_query(_query: str):
|
||||
await asyncio.sleep(0.08)
|
||||
return [{"result": 1}]
|
||||
|
||||
client.db.connect = AsyncMock(side_effect=_slow_connect)
|
||||
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
|
||||
client.db.query_raw = AsyncMock(side_effect=_slow_query)
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
with (
|
||||
pytest.raises(asyncio.TimeoutError),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=0.1)
|
||||
|
||||
|
||||
@@ -319,42 +335,154 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lightweight_reconnect_kills_engine_on_disconnect_failure(
|
||||
async def test_recreate_prisma_client_kills_old_engine_without_disconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Lightweight reconnect must kill the old engine PID when disconnect() fails."""
|
||||
"""recreate_prisma_client SIGTERMs the old engine PID directly rather than
|
||||
calling `disconnect()`, which blocks the asyncio event loop on the sync
|
||||
`subprocess.Popen.wait()` inside prisma-client-py — see issue #26191.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed"))
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
disconnect_mock = AsyncMock(
|
||||
side_effect=AssertionError("disconnect must not be called on reconnect path")
|
||||
)
|
||||
client.db._original_prisma.disconnect = disconnect_mock
|
||||
|
||||
with (
|
||||
patch.object(client, "_get_engine_pid", return_value=9999),
|
||||
patch("os.kill") as mock_kill,
|
||||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||||
patch.object(client.db, "_get_engine_pid", return_value=9999),
|
||||
patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill,
|
||||
patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
# Return a Prisma instance whose connect() is awaitable.
|
||||
fake_new_prisma = MagicMock()
|
||||
fake_new_prisma.connect = AsyncMock(return_value=None)
|
||||
with patch("prisma.Prisma", return_value=fake_new_prisma):
|
||||
await client.db.recreate_prisma_client("postgresql://test")
|
||||
|
||||
mock_kill.assert_any_call(9999, signal.SIGTERM)
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
disconnect_mock.assert_not_awaited()
|
||||
fake_new_prisma.connect.assert_awaited_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_generic_data: transport-reconnect-and-retry coverage (issue #25143)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
|
||||
async def test_get_generic_data_retries_on_transport_error_for_config_table(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Lightweight reconnect must NOT kill when disconnect() succeeds."""
|
||||
"""`get_generic_data(table_name="config")` self-heals on a transient
|
||||
`httpx.ReadError`: reconnect once, retry once, return the row.
|
||||
|
||||
Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry
|
||||
branch that 1.82.6 had on this method. `_update_config_from_db` fans out
|
||||
four concurrent `get_generic_data` calls, so a single transport flap used
|
||||
to surface as four `db_exceptions` alerts and a stale config window.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
|
||||
with patch("os.kill") as mock_kill:
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
|
||||
invocations: list[None] = []
|
||||
|
||||
mock_kill.assert_not_called()
|
||||
async def _flaky_find_first(**kwargs):
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("simulated transport blip")
|
||||
return expected_row
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
result = await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
assert result == expected_row
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
# The failure_handler telemetry side-effect must NOT fire on the first
|
||||
# transport blip — only if the post-retry call also fails. Drain the
|
||||
# event loop so any spuriously-spawned task would have run by now.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
|
||||
"""If reconnect itself does not succeed, propagate the original transport
|
||||
error and let the existing failure_handler / db_exceptions telemetry fire."""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(
|
||||
side_effect=httpx.ReadError("simulated transport blip")
|
||||
)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
# Failure telemetry IS expected here — the read genuinely failed.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _engine_confirmed_dead flag-reset bug (B2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Regression test for the flag-reset bug.
|
||||
|
||||
Before the fix, `_run_reconnect_cycle` cleared
|
||||
`self._engine_confirmed_dead = False` *before* awaiting
|
||||
`_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout,
|
||||
missing DATABASE_URL, recreate failure), the flag was left cleared and the
|
||||
next attempt could demote to the lightweight path even though the engine
|
||||
was genuinely dead.
|
||||
|
||||
The fix moves the reset into the success branch — the flag must stay True
|
||||
when heavy reconnect raises.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._engine_confirmed_dead = True
|
||||
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
|
||||
|
||||
# Make the heavy reconnect path raise.
|
||||
client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=RuntimeError("simulated heavy reconnect failure")
|
||||
)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
client._cleanup_engine_watcher = MagicMock()
|
||||
client._reap_all_zombies = MagicMock()
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
with pytest.raises(Exception):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
# The flag must STILL be True so the next attempt re-enters the heavy
|
||||
# branch instead of silently demoting to the lightweight path.
|
||||
assert client._engine_confirmed_dead is True
|
||||
|
||||
@@ -0,0 +1,611 @@
|
||||
"""
|
||||
Unit tests for workflow management endpoints (/v1/workflows/runs/*).
|
||||
Uses FastAPI TestClient with a mocked prisma_client.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.management_endpoints.workflow_management_endpoints import router
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_run(
|
||||
run_id: str = "run-1",
|
||||
session_id: str = "sess-1",
|
||||
workflow_type: str = "shin-builder",
|
||||
status: str = "pending",
|
||||
created_by: Any = "tok-test",
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.run_id = run_id
|
||||
obj.session_id = session_id
|
||||
obj.workflow_type = workflow_type
|
||||
obj.status = status
|
||||
obj.created_by = created_by
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
obj.updated_at = datetime.now(timezone.utc)
|
||||
obj.input = None
|
||||
obj.output = None
|
||||
obj.metadata = None
|
||||
return obj
|
||||
|
||||
|
||||
def _make_event(
|
||||
event_id: str = "evt-1",
|
||||
run_id: str = "run-1",
|
||||
event_type: str = "step.started",
|
||||
step_name: str = "grill",
|
||||
sequence_number: int = 0,
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.event_id = event_id
|
||||
obj.run_id = run_id
|
||||
obj.event_type = event_type
|
||||
obj.step_name = step_name
|
||||
obj.sequence_number = sequence_number
|
||||
obj.data = None
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
return obj
|
||||
|
||||
|
||||
def _make_message(
|
||||
message_id: str = "msg-1",
|
||||
run_id: str = "run-1",
|
||||
role: str = "user",
|
||||
content: str = "hello",
|
||||
sequence_number: int = 0,
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.message_id = message_id
|
||||
obj.run_id = run_id
|
||||
obj.role = role
|
||||
obj.content = content
|
||||
obj.sequence_number = sequence_number
|
||||
obj.session_id = None
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
return obj
|
||||
|
||||
|
||||
def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock:
|
||||
"""Build an async context-manager mock for prisma_client.db.tx()."""
|
||||
tx = MagicMock()
|
||||
tx.litellm_workflowevent = MagicMock()
|
||||
tx.litellm_workflowevent.create = AsyncMock(
|
||||
return_value=event_return or _make_event()
|
||||
)
|
||||
tx.litellm_workflowrun = MagicMock()
|
||||
tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run())
|
||||
tx.litellm_workflowmessage = MagicMock()
|
||||
tx.litellm_workflowmessage.create = AsyncMock(
|
||||
return_value=msg_return or _make_message()
|
||||
)
|
||||
tx.__aenter__ = AsyncMock(return_value=tx)
|
||||
tx.__aexit__ = AsyncMock(return_value=False)
|
||||
return tx
|
||||
|
||||
|
||||
def _make_prisma_client() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.db = MagicMock()
|
||||
client.db.litellm_workflowrun = MagicMock()
|
||||
client.db.litellm_workflowevent = MagicMock()
|
||||
client.db.litellm_workflowmessage = MagicMock()
|
||||
# default tx() returns a no-op transaction
|
||||
client.db.tx = MagicMock(return_value=_make_tx())
|
||||
return client
|
||||
|
||||
|
||||
def _make_app() -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
return app
|
||||
|
||||
|
||||
def _override_auth() -> Any:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-test", user_id="admin")
|
||||
auth.token = "tok-test"
|
||||
return auth
|
||||
|
||||
|
||||
def _override_auth_admin() -> Any:
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-master")
|
||||
auth.user_role = LitellmUserRoles.PROXY_ADMIN # type: ignore[assignment]
|
||||
return auth
|
||||
|
||||
|
||||
def _override_auth_user_with_token(token: str = "tok-abc") -> Any:
|
||||
"""Return a non-admin caller whose hashed token equals `token`."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-user", user_id="user-1")
|
||||
auth.token = token # override the computed hash with a predictable value
|
||||
return auth
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCreateWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_create_returns_run(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run())
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs",
|
||||
json={"workflow_type": "shin-builder"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
self._prisma.db.litellm_workflowrun.create.assert_awaited_once()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client", None)
|
||||
def test_create_500_when_no_db(self):
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs",
|
||||
json={"workflow_type": "shin-builder"},
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
|
||||
|
||||
class TestListWorkflowRuns:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_returns_runs(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(
|
||||
return_value=[_make_run()]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 1
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_filters_by_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs?status=running")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"]["status"] == "running"
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_filters_by_multiple_statuses(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs?status=running,paused")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"]["status"] == {"in": ["running", "paused"]}
|
||||
|
||||
|
||||
class TestGetWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_get_existing_run(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 200
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_get_missing_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/nonexistent")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestUpdateWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_update_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
updated = _make_run(status="completed")
|
||||
self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated)
|
||||
|
||||
resp = self.client.patch(
|
||||
"/v1/workflows/runs/run-1", json={"status": "completed"}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
self._prisma.db.litellm_workflowrun.update.assert_awaited_once()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_update_no_fields_returns_400(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
resp = self.client.patch("/v1/workflows/runs/run-1", json={})
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestAppendWorkflowEvent:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_event_updates_run_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
# _require_run check
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
tx = _make_tx(
|
||||
event_return=_make_event(), run_return=_make_run(status="running")
|
||||
)
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
# run status updated inside tx
|
||||
tx.litellm_workflowrun.update.assert_awaited_once()
|
||||
update_call = tx.litellm_workflowrun.update.call_args[1]
|
||||
assert update_call["data"]["status"] == "running"
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_event_no_status_update_for_unknown_type(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
tx = _make_tx(event_return=_make_event(event_type="custom.event"))
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "custom.event", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
# no status update inside tx for unknown event_type
|
||||
tx.litellm_workflowrun.update.assert_not_awaited()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_sequence_number_increments(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
existing = _make_event(sequence_number=4)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
|
||||
return_value=[existing]
|
||||
)
|
||||
tx = _make_tx(event_return=_make_event(sequence_number=5))
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "plan"},
|
||||
)
|
||||
create_call = tx.litellm_workflowevent.create.call_args[1]
|
||||
assert create_call["data"]["sequence_number"] == 5
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_unknown_run_id_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/nonexistent/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_sequence_collision_retries_and_succeeds(self, mock_pc):
|
||||
"""UniqueViolationError on first attempt triggers retry; second attempt succeeds."""
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# First tx raises UniqueViolationError; second succeeds.
|
||||
tx_fail = _make_tx()
|
||||
tx_fail.__aenter__ = AsyncMock(return_value=tx_fail)
|
||||
tx_fail.litellm_workflowevent.create = AsyncMock(
|
||||
side_effect=UniqueViolationError(
|
||||
{"user_facing_error": {"message": "unique"}}
|
||||
)
|
||||
)
|
||||
tx_fail.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
tx_ok = _make_tx(event_return=_make_event(sequence_number=1))
|
||||
|
||||
self._prisma.db.tx = MagicMock(side_effect=[tx_fail, tx_ok])
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
|
||||
class TestWorkflowMessages:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_message(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
|
||||
self._prisma.db.litellm_workflowmessage.create = AsyncMock(
|
||||
return_value=_make_message()
|
||||
)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/messages",
|
||||
json={"role": "user", "content": "fix the bug"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_message_unknown_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/nonexistent/messages",
|
||||
json={"role": "user", "content": "hello"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_messages_ordered(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_make_message(sequence_number=0),
|
||||
_make_message(sequence_number=1, role="assistant"),
|
||||
]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/messages")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 2
|
||||
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
|
||||
assert call_kwargs["order"] == {"sequence_number": "asc"}
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_messages_respects_limit(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
|
||||
assert call_kwargs["take"] == 25
|
||||
|
||||
|
||||
class TestListWorkflowEvents:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_ordered(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_make_event(sequence_number=0),
|
||||
_make_event(sequence_number=1),
|
||||
]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/events")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 2
|
||||
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
|
||||
assert call_kwargs["order"] == {"sequence_number": "asc"}
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_respects_limit(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
|
||||
assert call_kwargs["take"] == 10
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_unknown_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/nonexistent/events")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestTenantIsolation:
|
||||
"""Ownership enforcement: non-admin callers only see their own runs."""
|
||||
|
||||
def _make_app_with_auth(self, auth_fn):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = auth_fn
|
||||
return TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_create_stores_caller_token(self, mock_pc):
|
||||
token = "tok-owner"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.create = AsyncMock(
|
||||
return_value=_make_run(created_by=token)
|
||||
)
|
||||
|
||||
resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"})
|
||||
assert resp.status_code == 200
|
||||
create_call = self._prisma.db.litellm_workflowrun.create.call_args[1]
|
||||
assert create_call["data"]["created_by"] == token
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_list_scoped_to_caller_token(self, mock_pc):
|
||||
token = "tok-owner"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"].get("created_by") == token
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_admin_list_not_scoped(self, mock_pc):
|
||||
client = self._make_app_with_auth(_override_auth_admin)
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert "created_by" not in call_kwargs["where"]
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_other_users_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
# Run owned by a different key
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by="tok-other-owner")
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_null_owner_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=None)
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_update_null_owner_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=None)
|
||||
)
|
||||
self._prisma.db.litellm_workflowrun.update = AsyncMock(
|
||||
return_value=_make_run(status="completed")
|
||||
)
|
||||
|
||||
resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"})
|
||||
assert resp.status_code == 404
|
||||
self._prisma.db.litellm_workflowrun.update.assert_not_awaited()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_own_run_succeeds(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=token)
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 200
|
||||
@@ -218,6 +218,141 @@ class TestProxyBaseLLMRequestProcessing:
|
||||
headers_with_invalid
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_from_llm_response(self):
|
||||
"""
|
||||
Google native :generateContent uses this helper instead of base_process_llm_request;
|
||||
ensure x-litellm-* headers and callback hooks merge like the main proxy path.
|
||||
"""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
class _FakeGenaiResponse:
|
||||
_hidden_params = {
|
||||
"model_id": "deployment-model-id",
|
||||
"cache_key": "ck-test",
|
||||
"api_base": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"response_cost": 0.001,
|
||||
"additional_headers": {"llm_provider-ratelimit-requests": "1000"},
|
||||
}
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "call-id-test"
|
||||
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-ratelimit-remaining-requests": "999"}
|
||||
)
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_FakeGenaiResponse(),
|
||||
request_data={"model": "gemini/gemini-1.5-flash"},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="9.9.9",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-call-id"] == "call-id-test"
|
||||
assert headers["x-litellm-model-id"] == "deployment-model-id"
|
||||
assert headers["x-litellm-version"] == "9.9.9"
|
||||
assert headers["llm_provider-ratelimit-requests"] == "1000"
|
||||
assert headers["x-ratelimit-remaining-requests"] == "999"
|
||||
proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self):
|
||||
"""AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate."""
|
||||
|
||||
class _FakeStreamLike:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
_hidden_params = {
|
||||
"model_id": "stream-model-id",
|
||||
"api_base": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"cache_key": "",
|
||||
"response_cost": "",
|
||||
"additional_headers": {"llm_provider-x": "y"},
|
||||
}
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "cid-stream"
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_FakeStreamLike(),
|
||||
request_data={"model": "gemini/gemini-2.0-flash"},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="1.0.0",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-model-id"] == "stream-model-id"
|
||||
assert headers["x-litellm-model-api-base"] == (
|
||||
"https://generativelanguage.googleapis.com/v1beta"
|
||||
)
|
||||
assert headers["llm_provider-x"] == "y"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback(
|
||||
self,
|
||||
):
|
||||
"""When response has no _hidden_params, model_id can still come from litellm_metadata."""
|
||||
|
||||
class _BareResponse:
|
||||
pass
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "cid-meta"
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_BareResponse(),
|
||||
request_data={
|
||||
"model": "gemini/gemini-1.5-flash",
|
||||
"litellm_metadata": {"model_info": {"id": "meta-model-id"}},
|
||||
},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="1.0.0",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-model-id"] == "meta-model-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_with_stream_timeout_header(self):
|
||||
"""
|
||||
@@ -1158,13 +1293,6 @@ class TestCommonRequestProcessingHelpers:
|
||||
assert mock_tracer.trace.call_count == 4
|
||||
|
||||
# Verify that each call was made with the correct operation name
|
||||
expected_calls = [
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
]
|
||||
|
||||
actual_calls = mock_tracer.trace.call_args_list
|
||||
assert len(actual_calls) == 4
|
||||
|
||||
|
||||
@@ -2544,6 +2544,14 @@ class TestPriceDataReloadAPI:
|
||||
class TestPriceDataReloadIntegration:
|
||||
"""Integration tests for the complete price data reload feature"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flush_litellm_config_cache(self):
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
litellm_config_cache.flush_cache()
|
||||
yield
|
||||
litellm_config_cache.flush_cache()
|
||||
|
||||
@pytest.fixture
|
||||
def client_with_auth(self):
|
||||
"""Create a test client with authentication"""
|
||||
@@ -2601,6 +2609,7 @@ class TestPriceDataReloadIntegration:
|
||||
def test_distributed_reload_check_function(self):
|
||||
"""Test the _check_and_reload_model_cost_map function"""
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
proxy_config = ProxyConfig()
|
||||
|
||||
@@ -2609,14 +2618,19 @@ class TestPriceDataReloadIntegration:
|
||||
|
||||
# Test case 1: No config in database
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
# _check_and_reload_model_cost_map routes through get_config_param,
|
||||
# which calls prisma.get_generic_data on a cache miss.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=None)
|
||||
|
||||
# Should return early without reloading
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Test case 2: Config with interval but not time to reload
|
||||
litellm_config_cache.flush_cache()
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": False}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
|
||||
# Mock current time and last reload time
|
||||
with patch(
|
||||
@@ -2632,8 +2646,10 @@ class TestPriceDataReloadIntegration:
|
||||
asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma))
|
||||
|
||||
# Test case 3: Config with force reload
|
||||
litellm_config_cache.flush_cache()
|
||||
mock_config.param_value = {"interval_hours": 6, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
@@ -2675,6 +2691,8 @@ class TestPriceDataReloadIntegration:
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 24, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
# _check_and_reload_model_cost_map now reads through get_generic_data.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
original_model_cost = litellm.model_cost.copy()
|
||||
@@ -2770,6 +2788,8 @@ class TestPriceDataReloadIntegration:
|
||||
mock_config = MagicMock()
|
||||
mock_config.param_value = {"interval_hours": 12, "force_reload": True}
|
||||
mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config)
|
||||
# _check_and_reload_anthropic_beta_headers now reads through get_generic_data.
|
||||
mock_prisma.get_generic_data = AsyncMock(return_value=mock_config)
|
||||
mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None)
|
||||
|
||||
with patch(
|
||||
|
||||
@@ -29,8 +29,21 @@ from litellm.types.utils import (
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _use_local_model_cost_map(monkeypatch):
|
||||
original_model_cost = litellm.model_cost
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
litellm.get_model_info.cache_clear()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
class TestGPTImageCostCalculator:
|
||||
"""Test the OpenAI gpt-image-1 cost calculator"""
|
||||
"""Test the OpenAI gpt-image cost calculator"""
|
||||
|
||||
def test_gpt_image_1_cost_with_text_only(self):
|
||||
"""Test cost calculation with only text input tokens"""
|
||||
@@ -149,6 +162,44 @@ class TestGPTImageCostCalculator:
|
||||
|
||||
assert cost == 0.0
|
||||
|
||||
def test_gpt_image_2_cost_with_text_and_image_tokens(self):
|
||||
"""Test cost calculation for gpt-image-2 token pricing"""
|
||||
from litellm.llms.openai.image_generation.cost_calculator import cost_calculator
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=600,
|
||||
completion_tokens=5000,
|
||||
total_tokens=5600,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
text_tokens=100,
|
||||
image_tokens=500,
|
||||
),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=1000,
|
||||
image_tokens=4000,
|
||||
),
|
||||
)
|
||||
|
||||
image_response = ImageResponse(
|
||||
created=1234567890,
|
||||
data=[ImageObject(url="http://example.com/image.jpg")],
|
||||
)
|
||||
image_response.usage = usage
|
||||
|
||||
cost = cost_calculator(
|
||||
model="gpt-image-2",
|
||||
image_response=image_response,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
# GPT Image 2 pricing:
|
||||
# Text input: 100 * $5/1M = 0.0005
|
||||
# Image input: 500 * $8/1M = 0.004
|
||||
# Text output: 1000 * $10/1M = 0.01
|
||||
# Image output: 4000 * $30/1M = 0.12
|
||||
expected_cost = 0.0005 + 0.004 + 0.01 + 0.12
|
||||
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
|
||||
|
||||
|
||||
class TestGPTImageCostRouting:
|
||||
"""Test that gpt-image models are properly routed to the token-based calculator"""
|
||||
@@ -182,6 +233,33 @@ class TestGPTImageCostRouting:
|
||||
expected_cost = 0.0005 + 0.2
|
||||
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
|
||||
|
||||
def test_openai_gpt_image_2_routes_to_token_calculator(self):
|
||||
"""Test that OpenAI gpt-image-2 routes to token-based calculator"""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=5000,
|
||||
total_tokens=5100,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=5000),
|
||||
)
|
||||
|
||||
image_response = ImageResponse(
|
||||
created=1234567890,
|
||||
data=[ImageObject(url="http://example.com/image.jpg")],
|
||||
)
|
||||
image_response.usage = usage
|
||||
|
||||
cost = CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="gpt-image-2",
|
||||
completion_response=image_response,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
expected_cost = 0.0005 + 0.15
|
||||
assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"
|
||||
|
||||
def test_openai_dalle_routes_to_pixel_calculator(self):
|
||||
"""Test that OpenAI DALL-E still routes to pixel-based calculator"""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
|
||||
|
||||
@@ -32,6 +32,19 @@ from litellm.utils import (
|
||||
# Adds the parent directory to the system path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_model_cost_map(monkeypatch):
|
||||
original_model_cost = litellm.model_cost
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
litellm.get_model_info.cache_clear()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_check_provider_match_azure_ai_allows_openai_and_azure():
|
||||
"""
|
||||
Test that azure_ai provider can match openai and azure models.
|
||||
@@ -198,6 +211,72 @@ def test_get_optional_params_image_gen_filters_empty_values():
|
||||
assert optional_params == {}
|
||||
|
||||
|
||||
def test_gpt_image_provider_detection_covers_existing_family():
|
||||
for image_model in ("gpt-image-1", "gpt-image-1-mini", "gpt-image-1.5"):
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=image_model)
|
||||
|
||||
assert model == image_model
|
||||
assert custom_llm_provider == "openai"
|
||||
|
||||
|
||||
def test_gpt_image_2_provider_and_model_info(local_model_cost_map):
|
||||
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="gpt-image-2")
|
||||
|
||||
assert model == "gpt-image-2"
|
||||
assert custom_llm_provider == "openai"
|
||||
|
||||
model_info = litellm.get_model_info(model="gpt-image-2")
|
||||
assert model_info["litellm_provider"] == "openai"
|
||||
assert model_info["mode"] == "image_generation"
|
||||
assert model_info["input_cost_per_token"] == 5e-06
|
||||
assert model_info["input_cost_per_image_token"] == 8e-06
|
||||
assert model_info["output_cost_per_token"] == 1e-05
|
||||
assert model_info["output_cost_per_image_token"] == 3e-05
|
||||
assert (
|
||||
"/v1/images/generations"
|
||||
in litellm.model_cost["gpt-image-2"]["supported_endpoints"]
|
||||
)
|
||||
assert (
|
||||
"/v1/images/edits" in litellm.model_cost["gpt-image-2"]["supported_endpoints"]
|
||||
)
|
||||
assert model_info["supports_vision"] is True
|
||||
assert model_info["supports_pdf_input"] is True
|
||||
|
||||
|
||||
def test_gpt_image_2_snapshot_model_info(local_model_cost_map):
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model="gpt-image-2-2026-04-21"
|
||||
)
|
||||
|
||||
assert model == "gpt-image-2-2026-04-21"
|
||||
assert custom_llm_provider == "openai"
|
||||
|
||||
model_info = litellm.get_model_info(model="gpt-image-2-2026-04-21")
|
||||
assert model_info["litellm_provider"] == "openai"
|
||||
assert model_info["mode"] == "image_generation"
|
||||
assert model_info["output_cost_per_image_token"] == 3e-05
|
||||
|
||||
|
||||
def test_azure_gpt_image_2_model_info(local_model_cost_map):
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model="azure/gpt-image-2"
|
||||
)
|
||||
|
||||
assert model == "gpt-image-2"
|
||||
assert custom_llm_provider == "azure"
|
||||
|
||||
model_info = litellm.get_model_info(
|
||||
model="gpt-image-2", custom_llm_provider="azure"
|
||||
)
|
||||
assert model_info["litellm_provider"] == "azure"
|
||||
assert model_info["mode"] == "image_generation"
|
||||
assert model_info["input_cost_per_token"] == 5e-06
|
||||
assert model_info["input_cost_per_image_token"] == 8e-06
|
||||
assert model_info["output_cost_per_token"] == 1e-05
|
||||
assert model_info["output_cost_per_image_token"] == 3e-05
|
||||
|
||||
|
||||
def test_all_model_configs():
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import (
|
||||
VertexAIAi21Config,
|
||||
@@ -1179,7 +1258,7 @@ def test_get_model_info_shows_supports_computer_use():
|
||||
"model, custom_llm_provider",
|
||||
[
|
||||
("gpt-3.5-turbo", "openai"),
|
||||
("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"),
|
||||
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
|
||||
("gemini-2.5-pro", "vertex_ai"),
|
||||
],
|
||||
)
|
||||
@@ -1325,7 +1404,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/bedrock-claude-3-opus",
|
||||
"bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
),
|
||||
(
|
||||
@@ -1623,7 +1702,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/bedrock-claude-3-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Bedrock Claude 3 Opus via Converse API",
|
||||
),
|
||||
@@ -1710,7 +1789,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/staging-claude-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Staging Claude Opus",
|
||||
),
|
||||
@@ -1722,7 +1801,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/high-performance-claude",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"High-performance Claude deployment",
|
||||
),
|
||||
@@ -1860,7 +1939,7 @@ class TestProxyFunctionCalling:
|
||||
bedrock_models = [
|
||||
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
]
|
||||
|
||||
for model in bedrock_models:
|
||||
@@ -1892,7 +1971,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/bedrock-claude-3-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Bedrock Claude 3 Opus via Converse API",
|
||||
),
|
||||
@@ -1979,7 +2058,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/staging-claude-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Staging Claude Opus",
|
||||
),
|
||||
@@ -1991,7 +2070,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/high-performance-claude",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"High-performance Claude deployment",
|
||||
),
|
||||
@@ -2129,7 +2208,7 @@ class TestProxyFunctionCalling:
|
||||
bedrock_models = [
|
||||
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
]
|
||||
|
||||
for model in bedrock_models:
|
||||
@@ -2161,7 +2240,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/bedrock-claude-3-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Bedrock Claude 3 Opus via Converse API",
|
||||
),
|
||||
@@ -2248,7 +2327,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/staging-claude-opus",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"Staging Claude Opus",
|
||||
),
|
||||
@@ -2260,7 +2339,7 @@ class TestProxyFunctionCalling:
|
||||
),
|
||||
(
|
||||
"litellm_proxy/high-performance-claude",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
False,
|
||||
"High-performance Claude deployment",
|
||||
),
|
||||
@@ -2398,7 +2477,7 @@ class TestProxyFunctionCalling:
|
||||
bedrock_models = [
|
||||
"bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0",
|
||||
"bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
]
|
||||
|
||||
for model in bedrock_models:
|
||||
|
||||
@@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest):
|
||||
api_base="http://localhost:9380",
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
def test_transform_search_vector_store_response_not_implemented(self):
|
||||
|
||||
Reference in New Issue
Block a user