Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 1 addition & 19 deletions litellm/proxy/db/exception_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,25 +43,7 @@ def is_database_connection_error(e: Exception) -> bool:
):
return True
if isinstance(e, prisma.errors.PrismaError):
error_message = str(e).lower()
# Treat generic PrismaError as connection error only when its text
# clearly indicates transport/connectivity failure.
connection_keywords = (
"can't reach database server",
"cannot reach database server",
"can't connect",
"cannot connect",
"connection error",
"connection closed",
"timed out",
"timeout",
"connection refused",
"network is unreachable",
"no route to host",
"broken pipe",
)
if any(keyword in error_message for keyword in connection_keywords):
return True
return True
Comment on lines 45 to +46

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Existing unit test will fail

The test test_is_database_connection_error_non_connection_prisma_errors in tests/test_litellm/proxy/db/test_exception_handler.py:75-76 asserts that PrismaError(), DataError, UniqueViolationError, ForeignKeyViolationError, MissingRequiredValueError, RawQueryError, TableNotFoundError, and RecordNotFoundError all return False from is_database_connection_error. With this change, all of them will return True, causing 8 test failures.

This test file was not updated as part of this PR. It needs to be updated to match the new intended behavior (either by asserting True for those cases, or removing those test cases).

Context Used: Rule from dashboard - What: Ensure that any PR claiming to fix an issue includes evidence that the issue is resolved, such... (source)

Comment on lines 45 to +46

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-connection errors misclassified as connection errors

This broadening causes unintended side effects in call sites that use is_database_connection_error for purposes other than allow_requests_on_db_unavailable:

  1. access_group_endpoints.py:333: A RecordNotFoundError during delete_access_group will now match is_database_connection_error and return 503 Service Unavailable instead of falling through to the P2025/not-found check on line 338, which would correctly return 404.

  2. proxy/utils.py:4113: The DB health watchdog uses is_database_connection_error to decide whether to trigger a reconnect. With this change, any PrismaError (e.g. a query syntax error or unique violation from the SELECT 1 probe β€” unlikely but possible in edge cases) would trigger unnecessary reconnection attempts.

  3. auth/auth_checks.py:2003: Same pattern β€” a UniqueViolationError or DataError during key lookup would trigger a DB reconnect attempt when no connectivity issue actually exists.

Consider either:

  • Keeping the broad match only for the handle_db_exception / allow_requests_on_db_unavailable path (separate method), or
  • At minimum, reordering the check in access_group_endpoints.py so the P2025 check runs before is_database_connection_error.

if isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection:
return True
return False
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
"user": "",
"team_id": "",
"organization_id": "",
"metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}}",
"metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}",
"cache_key": "Cache OFF",
"spend": 0.00022500000000000002,
"total_tokens": 30,
Expand Down
165 changes: 126 additions & 39 deletions tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,19 +176,49 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_streaming(s
"""
Test that the Bedrock Knowledge Base Hook works with streaming and returns search_results in chunks.
"""

# Init client
# litellm._turn_on_debug()
async_client = AsyncHTTPHandler()
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
vector_store_ids = [
"T37J8R4WTM"
],
stream=True,
client=async_client
)

async def mock_anthropic_aiter_lines():
lines = [
'event: message_start',
'data: {"type":"message_start","message":{"id":"msg_01ABC","type":"message","role":"assistant","content":[],"model":"claude-3-5-haiku-latest","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}',
'',
'event: content_block_start',
'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
'',
'event: content_block_delta',
'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"LiteLLM is a library."}}',
'',
'event: content_block_stop',
'data: {"type":"content_block_stop","index":0}',
'',
'event: message_delta',
'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}',
'',
'event: message_stop',
'data: {"type":"message_stop"}',
]
for line in lines:
yield line

mock_streaming_response = Mock()
mock_streaming_response.status_code = 200
mock_streaming_response.headers = {}
mock_streaming_response.aiter_lines = mock_anthropic_aiter_lines

with patch.object(async_client, "post", new=AsyncMock(return_value=mock_streaming_response)):
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
vector_store_ids = [
"T37J8R4WTM"
],
stream=True,
client=async_client
)

# Collect chunks
chunks = []
Expand Down Expand Up @@ -228,20 +258,40 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools(
"""
Test that the Bedrock Knowledge Base Hook works when making a real llm api call
"""

# Init client
litellm._turn_on_debug()
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
max_tokens=10,
tools=[
{
"type": "file_search",
"vector_store_ids": ["T37J8R4WTM"]
}
],
)
async_client = AsyncHTTPHandler()

mock_llm_response = Mock()
mock_llm_response.status_code = 200
mock_llm_response.headers = {"content-type": "application/json"}
mock_llm_response_json = {
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library."}],
"model": "claude-3-5-haiku-latest",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
mock_llm_response.text = json.dumps(mock_llm_response_json)
mock_llm_response.json = Mock(return_value=mock_llm_response_json)

with patch.object(async_client, "post", new=AsyncMock(return_value=mock_llm_response)):
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
max_tokens=10,
tools=[
{
"type": "file_search",
"vector_store_ids": ["T37J8R4WTM"]
}
],
client=async_client,
)
assert response is not None

@pytest.mark.asyncio
Expand All @@ -253,23 +303,42 @@ async def test_e2e_bedrock_knowledgebase_retrieval_with_llm_api_call_with_tools_
In this case we filter for a non-existent user_id, which should return no results.
"""
litellm._turn_on_debug()

response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
max_tokens=10,
tools=[
{
"type": "file_search",
"vector_store_ids": ["T37J8R4WTM"],
"filters": {
"key": "user_id",
"value": "fake-user-id",
"operator": "eq"
async_client = AsyncHTTPHandler()

mock_llm_response = Mock()
mock_llm_response.status_code = 200
mock_llm_response.headers = {"content-type": "application/json"}
mock_llm_response_json = {
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library."}],
"model": "claude-3-5-haiku-latest",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
mock_llm_response.text = json.dumps(mock_llm_response_json)
mock_llm_response.json = Mock(return_value=mock_llm_response_json)

with patch.object(async_client, "post", new=AsyncMock(return_value=mock_llm_response)):
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
max_tokens=10,
tools=[
{
"type": "file_search",
"vector_store_ids": ["T37J8R4WTM"],
"filters": {
"key": "user_id",
"value": "fake-user-id",
"operator": "eq"
}
}
}
],
)
],
client=async_client,
)

# Verify response is not None
assert response is not None
Expand Down Expand Up @@ -344,11 +413,28 @@ async def fake_async_vector_store_search_handler(
],
)

async_client = AsyncHTTPHandler()
mock_llm_response = Mock()
mock_llm_response.status_code = 200
mock_llm_response.headers = {"content-type": "application/json"}
mock_llm_response_json = {
"id": "msg_01ABC123",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "LiteLLM is a library."}],
"model": "claude-3-5-haiku-latest",
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 5},
}
mock_llm_response.text = json.dumps(mock_llm_response_json)
mock_llm_response.json = Mock(return_value=mock_llm_response_json)

with patch.object(
litellm.vector_stores.main.base_llm_http_handler,
"async_vector_store_search_handler",
new=AsyncMock(side_effect=fake_async_vector_store_search_handler),
):
), patch.object(async_client, "post", new=AsyncMock(return_value=mock_llm_response)):
response = await litellm.acompletion(
model="anthropic/claude-3-5-haiku-latest",
messages=[{"role": "user", "content": "what is litellm?"}],
Expand All @@ -364,6 +450,7 @@ async def fake_async_vector_store_search_handler(
},
}
],
client=async_client,
)

assert response is not None
Expand Down
Loading