From 498c4254ead47876814b8a4f9b2fce485b943a37 Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sun, 7 Dec 2025 20:53:51 +0800 Subject: [PATCH 001/330] fix: Return 403 exception when calling GET responses api --- litellm/proxy/auth/auth_checks.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fc79a4d3591..309bd577606 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -402,13 +402,14 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: - user_route: str - the route the user is trying to call - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. """ + from starlette.routing import compile_path for allowed_route in allowed_routes: - if ( - allowed_route in LiteLLMRoutes.__members__ - and user_route in LiteLLMRoutes[allowed_route].value - ): - return True + if allowed_route in LiteLLMRoutes.__members__: + for template in LiteLLMRoutes[allowed_route].value: + regex, _, _ = compile_path(template) + if regex.match(user_route): + return True elif allowed_route == user_route: return True return False From 4ab58619ad56d93eb4add67bfd6064c50f449fa1 Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sun, 14 Dec 2025 17:55:13 +0800 Subject: [PATCH 002/330] fix: added new step into rotate master key function for processing credentials table --- .../proxy/credential_endpoints/endpoints.py | 9 ++--- .../key_management_endpoints.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 647abb73648..9f228bb1184 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -21,11 +21,11 @@ router = APIRouter() class CredentialHelperUtils: @staticmethod - def encrypt_credential_values(credential: CredentialItem) -> CredentialItem: + def encrypt_credential_values(credential: CredentialItem, new_encryption_key: Optional[str] = None) -> CredentialItem: """Encrypt values in credential.credential_values and add to DB""" encrypted_credential_values = {} for key, value in (credential.credential_values or {}).items(): - encrypted_credential_values[key] = encrypt_value_helper(value) + encrypted_credential_values[key] = encrypt_value_helper(value, new_encryption_key) # Return a new object to avoid mutating the caller's credential, which # is kept in memory and should remain unencrypted. @@ -246,7 +246,7 @@ async def delete_credential( def update_db_credential( - db_credential: CredentialItem, updated_patch: CredentialItem + db_credential: CredentialItem, updated_patch: CredentialItem, new_encryption_key: Optional[str] = None ) -> CredentialItem: """ Update a credential in the DB. @@ -258,7 +258,8 @@ def update_db_credential( ) encrypted_credential = CredentialHelperUtils.encrypt_credential_values( - updated_patch + updated_patch, + new_encryption_key, ) # update model name if encrypted_credential.credential_name: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index da44bda791d..8ea3122ce01 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2539,6 +2539,40 @@ async def _rotate_master_key( new_master_key=new_master_key, ) + # 5. process credentials table + try: + credentials = await prisma_client.db.litellm_credentialstable.find_many() + except Exception: + credentials = None + if credentials: + from litellm.proxy.credential_endpoints.endpoints import update_db_credential + + for cred in credentials: + try: + decrypted_cred = proxy_config.decrypt_credentials(cred) + encrypted_cred = update_db_credential( + db_credential=cred, + updated_patch=decrypted_cred, + new_encryption_key=new_master_key, + ) + credential_object_jsonified = jsonify_object(encrypted_cred.model_dump()) + await prisma_client.db.litellm_credentialstable.update( + where={"credential_name": cred.credential_name}, + data={ + **credential_object_jsonified, + "updated_by": user_api_key_dict.user_id, + }, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to re-encrypt credential {cred.credential_name}: {str(e)}" + ) + # Continue with next credential instead of failing entire rotation + continue + verbose_proxy_logger.debug( + f"Successfully re-encrypted {len(credentials)} credentials with new master key" + ) + def get_new_token(data: Optional[RegenerateKeyRequest]) -> str: if data and data.new_key is not None: From 7b6a00d3df82352ae1cdb76f6e7a5c6ac632b7e6 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 16 Dec 2025 13:06:33 -0300 Subject: [PATCH 003/330] fix: remove deprecated Groq models and update model registry - Remove 20 deprecated/unavailable Groq models from registry - Add groq/meta-llama/llama-guard-4-12b (new safety model) - Add supports_vision to Llama 4 models (maverick, scout) - Update Groq documentation with current model list - Clean up test file references to deprecated models Fixes #18043 --- docs/my-website/docs/providers/groq.md | 31 +-- ...odel_prices_and_context_window_backup.json | 240 +----------------- model_prices_and_context_window.json | 240 +----------------- tests/test_litellm/test_utils.py | 2 - 4 files changed, 24 insertions(+), 489 deletions(-) diff --git a/docs/my-website/docs/providers/groq.md b/docs/my-website/docs/providers/groq.md index ebed31f720f..55c222635d2 100644 --- a/docs/my-website/docs/providers/groq.md +++ b/docs/my-website/docs/providers/groq.md @@ -150,15 +150,15 @@ We support ALL Groq models, just set `groq/` as a prefix when sending completion | Model Name | Usage | |--------------------|---------------------------------------------------------| -| llama-3.1-8b-instant | `completion(model="groq/llama-3.1-8b-instant", messages)` | -| llama-3.1-70b-versatile | `completion(model="groq/llama-3.1-70b-versatile", messages)` | -| llama3-8b-8192 | `completion(model="groq/llama3-8b-8192", messages)` | -| llama3-70b-8192 | `completion(model="groq/llama3-70b-8192", messages)` | -| llama2-70b-4096 | `completion(model="groq/llama2-70b-4096", messages)` | -| mixtral-8x7b-32768 | `completion(model="groq/mixtral-8x7b-32768", messages)` | -| gemma-7b-it | `completion(model="groq/gemma-7b-it", messages)` | -| moonshotai/kimi-k2-instruct | `completion(model="groq/moonshotai/kimi-k2-instruct", messages)` | -| qwen3-32b | `completion(model="groq/qwen/qwen3-32b", messages)` | +| llama-3.3-70b-versatile | `completion(model="groq/llama-3.3-70b-versatile", messages)` | +| llama-3.1-8b-instant | `completion(model="groq/llama-3.1-8b-instant", messages)` | +| meta-llama/llama-4-scout-17b-16e-instruct | `completion(model="groq/meta-llama/llama-4-scout-17b-16e-instruct", messages)` | +| meta-llama/llama-4-maverick-17b-128e-instruct | `completion(model="groq/meta-llama/llama-4-maverick-17b-128e-instruct", messages)` | +| meta-llama/llama-guard-4-12b | `completion(model="groq/meta-llama/llama-guard-4-12b", messages)` | +| qwen/qwen3-32b | `completion(model="groq/qwen/qwen3-32b", messages)` | +| moonshotai/kimi-k2-instruct-0905 | `completion(model="groq/moonshotai/kimi-k2-instruct-0905", messages)` | +| openai/gpt-oss-120b | `completion(model="groq/openai/gpt-oss-120b", messages)` | +| openai/gpt-oss-20b | `completion(model="groq/openai/gpt-oss-20b", messages)` | ## Groq - Tool / Function Calling Example @@ -261,31 +261,28 @@ if tool_calls: print("second response\n", second_response) ``` -## Groq - Vision Example +## Groq - Vision Example -Select Groq models support vision. Check out their [model list](https://console.groq.com/docs/vision) for more details. +Groq's Llama 4 models support vision. Check out their [model list](https://console.groq.com/docs/vision) for more details. ```python -from litellm import completion - -import os +import os from litellm import completion os.environ["GROQ_API_KEY"] = "your-api-key" -# openai call response = completion( - model = "groq/llama-3.2-11b-vision-preview", + model = "groq/meta-llama/llama-4-scout-17b-16e-instruct", messages=[ { "role": "user", "content": [ { "type": "text", - "text": "What’s in this image?" + "text": "What's in this image?" }, { "type": "image_url", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2af548ce07e..5754899846c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17509,75 +17509,6 @@ "supports_response_schema": true, "supports_vision": true }, - "groq/deepseek-r1-distill-llama-70b": { - "input_cost_per_token": 7.5e-07, - "litellm_provider": "groq", - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 9.9e-07, - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/distil-whisper-large-v3-en": { - "input_cost_per_second": 5.56e-06, - "litellm_provider": "groq", - "mode": "audio_transcription", - "output_cost_per_second": 0.0 - }, - "groq/gemma-7b-it": { - "deprecation_date": "2024-12-18", - "input_cost_per_token": 7e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/gemma2-9b-it": { - "input_cost_per_token": 2e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2e-07, - "supports_function_calling": false, - "supports_response_schema": false, - "supports_tool_choice": false - }, - "groq/llama-3.1-405b-reasoning": { - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.1-70b-versatile": { - "deprecation_date": "2025-01-24", - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, "groq/llama-3.1-8b-instant": { "input_cost_per_token": 5e-08, "litellm_provider": "groq", @@ -17590,97 +17521,6 @@ "supports_response_schema": false, "supports_tool_choice": true }, - "groq/llama-3.2-11b-text-preview": { - "deprecation_date": "2024-10-28", - "input_cost_per_token": 1.8e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-11b-vision-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 1.8e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "groq/llama-3.2-1b-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 4e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-3b-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 6e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-90b-text-preview": { - "deprecation_date": "2024-11-25", - "input_cost_per_token": 9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-90b-vision-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "groq/llama-3.3-70b-specdec": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9.9e-07, - "supports_tool_choice": true - }, "groq/llama-3.3-70b-versatile": { "input_cost_per_token": 5.9e-07, "litellm_provider": "groq", @@ -17693,7 +17533,7 @@ "supports_response_schema": false, "supports_tool_choice": true }, - "groq/llama-guard-3-8b": { + "groq/meta-llama/llama-guard-4-12b": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", "max_input_tokens": 8192, @@ -17702,44 +17542,6 @@ "mode": "chat", "output_cost_per_token": 2e-07 }, - "groq/llama2-70b-4096": { - "input_cost_per_token": 7e-07, - "litellm_provider": "groq", - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama3-groq-70b-8192-tool-use-preview": { - "deprecation_date": "2025-01-06", - "input_cost_per_token": 8.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 8.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama3-groq-8b-8192-tool-use-preview": { - "deprecation_date": "2025-01-06", - "input_cost_per_token": 1.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, "groq/meta-llama/llama-4-maverick-17b-128e-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", @@ -17750,7 +17552,8 @@ "output_cost_per_token": 6e-07, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "groq/meta-llama/llama-4-scout-17b-16e-instruct": { "input_cost_per_token": 1.1e-07, @@ -17762,41 +17565,8 @@ "output_cost_per_token": 3.4e-07, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true - }, - "groq/mistral-saba-24b": { - "input_cost_per_token": 7.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 32000, - "max_output_tokens": 32000, - "max_tokens": 32000, - "mode": "chat", - "output_cost_per_token": 7.9e-07 - }, - "groq/mixtral-8x7b-32768": { - "deprecation_date": "2025-03-20", - "input_cost_per_token": 2.4e-07, - "litellm_provider": "groq", - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 2.4e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/moonshotai/kimi-k2-instruct": { - "input_cost_per_token": 1e-06, - "litellm_provider": "groq", - "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 3e-06, - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "groq/moonshotai/kimi-k2-instruct-0905": { "input_cost_per_token": 1e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2af548ce07e..5754899846c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17509,75 +17509,6 @@ "supports_response_schema": true, "supports_vision": true }, - "groq/deepseek-r1-distill-llama-70b": { - "input_cost_per_token": 7.5e-07, - "litellm_provider": "groq", - "max_input_tokens": 128000, - "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", - "output_cost_per_token": 9.9e-07, - "supports_function_calling": true, - "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/distil-whisper-large-v3-en": { - "input_cost_per_second": 5.56e-06, - "litellm_provider": "groq", - "mode": "audio_transcription", - "output_cost_per_second": 0.0 - }, - "groq/gemma-7b-it": { - "deprecation_date": "2024-12-18", - "input_cost_per_token": 7e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/gemma2-9b-it": { - "input_cost_per_token": 2e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 2e-07, - "supports_function_calling": false, - "supports_response_schema": false, - "supports_tool_choice": false - }, - "groq/llama-3.1-405b-reasoning": { - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.1-70b-versatile": { - "deprecation_date": "2025-01-24", - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 7.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, "groq/llama-3.1-8b-instant": { "input_cost_per_token": 5e-08, "litellm_provider": "groq", @@ -17590,97 +17521,6 @@ "supports_response_schema": false, "supports_tool_choice": true }, - "groq/llama-3.2-11b-text-preview": { - "deprecation_date": "2024-10-28", - "input_cost_per_token": 1.8e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-11b-vision-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 1.8e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "groq/llama-3.2-1b-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 4e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 4e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-3b-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 6e-08, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-90b-text-preview": { - "deprecation_date": "2024-11-25", - "input_cost_per_token": 9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama-3.2-90b-vision-preview": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "groq/llama-3.3-70b-specdec": { - "deprecation_date": "2025-04-14", - "input_cost_per_token": 5.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 9.9e-07, - "supports_tool_choice": true - }, "groq/llama-3.3-70b-versatile": { "input_cost_per_token": 5.9e-07, "litellm_provider": "groq", @@ -17693,7 +17533,7 @@ "supports_response_schema": false, "supports_tool_choice": true }, - "groq/llama-guard-3-8b": { + "groq/meta-llama/llama-guard-4-12b": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", "max_input_tokens": 8192, @@ -17702,44 +17542,6 @@ "mode": "chat", "output_cost_per_token": 2e-07 }, - "groq/llama2-70b-4096": { - "input_cost_per_token": 7e-07, - "litellm_provider": "groq", - "max_input_tokens": 4096, - "max_output_tokens": 4096, - "max_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 8e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama3-groq-70b-8192-tool-use-preview": { - "deprecation_date": "2025-01-06", - "input_cost_per_token": 8.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 8.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/llama3-groq-8b-8192-tool-use-preview": { - "deprecation_date": "2025-01-06", - "input_cost_per_token": 1.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 8192, - "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", - "output_cost_per_token": 1.9e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, "groq/meta-llama/llama-4-maverick-17b-128e-instruct": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", @@ -17750,7 +17552,8 @@ "output_cost_per_token": 6e-07, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "groq/meta-llama/llama-4-scout-17b-16e-instruct": { "input_cost_per_token": 1.1e-07, @@ -17762,41 +17565,8 @@ "output_cost_per_token": 3.4e-07, "supports_function_calling": true, "supports_response_schema": true, - "supports_tool_choice": true - }, - "groq/mistral-saba-24b": { - "input_cost_per_token": 7.9e-07, - "litellm_provider": "groq", - "max_input_tokens": 32000, - "max_output_tokens": 32000, - "max_tokens": 32000, - "mode": "chat", - "output_cost_per_token": 7.9e-07 - }, - "groq/mixtral-8x7b-32768": { - "deprecation_date": "2025-03-20", - "input_cost_per_token": 2.4e-07, - "litellm_provider": "groq", - "max_input_tokens": 32768, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 2.4e-07, - "supports_function_calling": true, - "supports_response_schema": false, - "supports_tool_choice": true - }, - "groq/moonshotai/kimi-k2-instruct": { - "input_cost_per_token": 1e-06, - "litellm_provider": "groq", - "max_input_tokens": 131072, - "max_output_tokens": 16384, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 3e-06, - "supports_function_calling": true, - "supports_response_schema": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_vision": true }, "groq/moonshotai/kimi-k2-instruct-0905": { "input_cost_per_token": 1e-06, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 2bd94488ba2..a3dc6f5085d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -870,8 +870,6 @@ SKIP_MODELS = [ "jamba", "deepinfra", "mistral.", - "groq/llama-guard-3-8b", - "groq/gemma2-9b-it", ] # Bedrock models to block - organized by type From 9274860aa28c5d69d696077b92283dfcad90e0a5 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 14:54:26 -0800 Subject: [PATCH 004/330] Base commit --- litellm/proxy/management_endpoints/ui_sso.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d1db21a2706..5094fc5de97 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -594,7 +594,7 @@ def _build_sso_user_update_data( user_id: Optional[str], ) -> dict: """ - Build the update data dictionary for SSO user upsert. + Build the update data dictionary for SSO user upsert Args: result: The SSO response containing user information From 42d7d757a3bb1e5da02ed53c494ce6abca2f7e61 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 15:48:26 -0800 Subject: [PATCH 005/330] Adding role mappings to SSOConfig DB --- litellm/proxy/management_endpoints/ui_sso.py | 2 +- litellm/proxy/proxy_server.py | 1 + .../proxy_setting_endpoints.py | 12 ++ .../proxy/management_endpoints/ui_sso.py | 34 ++++- .../test_proxy_setting_endpoints.py | 135 ++++++++++++++++++ 5 files changed, 182 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 5094fc5de97..d1db21a2706 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -594,7 +594,7 @@ def _build_sso_user_update_data( user_id: Optional[str], ) -> dict: """ - Build the update data dictionary for SSO user upsert + Build the update data dictionary for SSO user upsert. Args: result: The SSO response containing user information diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8fdea95d7ad..dfadab1d531 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3578,6 +3578,7 @@ class ProxyConfig: ) if sso_settings is not None: # Capitalize all keys in sso_settings dictionary + sso_settings.sso_settings.pop("role_mappings", None) uppercase_sso_settings = { key.upper(): value for key, value in sso_settings.sso_settings.items() diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 9c99b625e9f..d9a41d38b22 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -433,10 +433,21 @@ async def get_sso_settings(): if sso_db_record and sso_db_record.sso_settings: # Load settings from database sso_settings_dict = dict(sso_db_record.sso_settings) + + # Extract role_mappings before removing it (it's a dict, not an env variable) + role_mappings_data = sso_settings_dict.pop("role_mappings", None) + role_mappings = None + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables(environment_variables=sso_settings_dict) # Build SSO config with database values or environment fallback + sso_config = SSOConfig( google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), @@ -451,6 +462,7 @@ async def get_sso_settings(): proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), user_email=decrypted_sso_settings_dict.get("user_email"), ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), + role_mappings=role_mappings, ) # Get the schema for UI display diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 820b0164400..187d8c97c05 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -1,10 +1,12 @@ -from typing import List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from pydantic import Field from typing_extensions import TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase +from litellm.proxy._types import LitellmUserRoles + class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): """ @@ -60,6 +62,30 @@ class AccessControl_UI_AccessMode(LiteLLMPydanticObjectBase): sso_group_jwt_field: str +class RoleMappings(LiteLLMPydanticObjectBase): + """ + Configuration for mapping SSO groups to LiteLLM roles. + + The system will look at the group_claim field in the SSO token to determine + which role to assign the user based on the roles mapping. + """ + + provider: str = Field( + description="SSO Provider name (e.g., 'google', 'microsoft', 'generic')" + ) + group_claim: str = Field( + description="The field name in the SSO token that contains the groups array (e.g., 'groups', 'roles')" + ) + default_role: Optional[LitellmUserRoles] = Field( + default=None, + description="Default role to assign if user's groups don't match any role mappings. Must be a valid LitellmUserRoles value (e.g., 'proxy_admin', 'internal_user', 'proxy_admin_viewer')" + ) + roles: Dict[LitellmUserRoles, List[str]] = Field( + default_factory=dict, + description="Mapping of LiteLLM role names to arrays of SSO group names. Example: {'proxy_admin': ['group-1', 'group-2'], 'proxy_admin_viewer': ['group-3']}" + ) + + class SSOConfig(LiteLLMPydanticObjectBase): """ Configuration for SSO environment variables and settings @@ -127,6 +153,12 @@ class SSOConfig(LiteLLMPydanticObjectBase): description="Access mode for the UI", ) + # Role Mappings + role_mappings: Optional[RoleMappings] = Field( + default=None, + description="Configuration for mapping SSO groups to LiteLLM roles based on group claims in the SSO token", + ) + class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): """ diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index d3c99151195..8fdfd6897a8 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -290,6 +290,10 @@ class TestProxySettingEndpoints: assert "google_client_id" in data["field_schema"]["properties"] assert "description" in data["field_schema"]["properties"]["google_client_id"] + # Verify role_mappings is present in response (can be None if not set) + assert "role_mappings" in values + assert values["role_mappings"] is None + # Verify find_unique was called with correct parameters mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once() call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args @@ -863,6 +867,10 @@ class TestProxySettingEndpoints: assert values["google_client_secret"] == "decrypted_google_secret" assert values["microsoft_client_id"] == "decrypted_microsoft_id" assert values["proxy_base_url"] == "https://decrypted.example.com" + + # Verify role_mappings is present in response (can be None if not set) + assert "role_mappings" in values + assert values["role_mappings"] is None def test_update_sso_settings_to_database(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating SSO settings saves to the dedicated database table""" @@ -1062,6 +1070,7 @@ class TestProxySettingEndpoints: assert values.get("google_client_id") is None assert values.get("google_client_secret") is None assert values.get("microsoft_client_id") is None + assert values.get("role_mappings") is None def test_update_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch): """Test updating SSO settings when database is not connected""" @@ -1088,3 +1097,129 @@ class TestProxySettingEndpoints: data = response.json() assert "error" in data["detail"] assert "Database not connected" in data["detail"]["error"] + + def test_get_sso_settings_with_role_mappings(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting SSO settings when role_mappings is present in database""" + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + # Mock the prisma client with database record containing role_mappings + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.sso_settings = { + "google_client_id": "test_google_client_id", + "role_mappings": { + "provider": "google", + "group_claim": "groups", + "default_role": LitellmUserRoles.INTERNAL_USER, + "roles": { + LitellmUserRoles.PROXY_ADMIN: ["admin-group"], + }, + }, + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock decryption to return the values as-is (role_mappings should not be passed to decryption) + from litellm.proxy.proxy_server import proxy_config + def mock_decrypt(environment_variables): + # role_mappings should not be in environment_variables since it's extracted before decryption + assert "role_mappings" not in environment_variables + return environment_variables + + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + # Verify role_mappings is returned correctly + values = data["values"] + assert "role_mappings" in values + assert values["role_mappings"] is not None + assert values["role_mappings"]["provider"] == "google" + assert values["role_mappings"]["group_claim"] == "groups" + assert values["role_mappings"]["default_role"] == LitellmUserRoles.INTERNAL_USER + assert values["role_mappings"]["roles"][LitellmUserRoles.PROXY_ADMIN] == ["admin-group"] + + def test_role_mappings_stored_and_retrieved(self, mock_proxy_config, mock_auth, monkeypatch): + """Test that role_mappings is properly stored and retrieved from SSO settings""" + import json + from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + mock_prisma.db.litellm_config = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) + + # SSO settings with role_mappings + role_mappings_data = { + "provider": "google", + "group_claim": "groups", + "default_role": LitellmUserRoles.INTERNAL_USER, + "roles": { + LitellmUserRoles.PROXY_ADMIN: ["admin-group"], + LitellmUserRoles.INTERNAL_USER: ["user-group"], + }, + } + + new_sso_settings = { + "google_client_id": "test_google_id", + "role_mappings": role_mappings_data, + } + + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert "role_mappings" in data["settings"] + + # Verify role_mappings structure in response + returned_role_mappings = data["settings"]["role_mappings"] + assert returned_role_mappings["provider"] == "google" + assert returned_role_mappings["group_claim"] == "groups" + assert returned_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER + assert returned_role_mappings["roles"][LitellmUserRoles.PROXY_ADMIN] == ["admin-group"] + + # Verify upsert was called with role_mappings in the data + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + create_data = call_args.kwargs["data"]["create"] + stored_sso_settings = json.loads(create_data["sso_settings"]) + assert "role_mappings" in stored_sso_settings + assert stored_sso_settings["role_mappings"]["provider"] == "google" + + # Now test retrieving role_mappings + mock_db_record = MagicMock() + mock_db_record.sso_settings = stored_sso_settings + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables + ) + + get_response = client.get("/get/sso_settings") + assert get_response.status_code == 200 + get_data = get_response.json() + + # Verify role_mappings is returned correctly + assert "role_mappings" in get_data["values"] + retrieved_role_mappings = get_data["values"]["role_mappings"] + assert retrieved_role_mappings is not None + assert retrieved_role_mappings["provider"] == "google" + assert retrieved_role_mappings["group_claim"] == "groups" + assert retrieved_role_mappings["default_role"] == LitellmUserRoles.INTERNAL_USER From 58330f852d82bd42f02f2e920cdb14ae0b927024 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 16 Dec 2025 17:36:08 -0800 Subject: [PATCH 006/330] WIP waiting for okta --- litellm/proxy/management_endpoints/ui_sso.py | 134 ++++++++++++++++++- 1 file changed, 127 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d1db21a2706..7a404bcda81 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -85,6 +85,58 @@ else: router = APIRouter() +def determine_role_from_groups( + user_groups: List[str], + role_mappings: "RoleMappings", +) -> Optional[LitellmUserRoles]: + """ + Determine the highest privilege role for a user based on their groups. + + Role hierarchy (highest to lowest): + - proxy_admin + - proxy_admin_viewer + - internal_user + - internal_user_viewer + + Args: + user_groups: List of group names from the SSO token + role_mappings: RoleMappings configuration object + + Returns: + The highest privilege role found, or default_role if no matches, or None + """ + if not role_mappings.roles: + # No role mappings configured, return default_role + return role_mappings.default_role + + # Role hierarchy (highest to lowest) + role_hierarchy = [ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + + # Convert user_groups to a set for efficient lookup + user_groups_set = set(user_groups) if isinstance(user_groups, list) else set() + + # Find the highest privilege role the user belongs to + for role in role_hierarchy: + if role in role_mappings.roles: + role_groups = role_mappings.roles[role] + if isinstance(role_groups, list) and user_groups_set.intersection(set(role_groups)): + verbose_proxy_logger.debug( + f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}" + ) + return role + + # No matching groups found, return default_role + verbose_proxy_logger.debug( + f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}" + ) + return role_mappings.default_role + + def process_sso_jwt_access_token( access_token_str: Optional[str], sso_jwt_handler: Optional[JWTHandler], @@ -243,6 +295,7 @@ def generic_response_convertor( response, jwt_handler: JWTHandler, sso_jwt_handler: Optional[JWTHandler] = None, + role_mappings: Optional["RoleMappings"] = None, ) -> CustomOpenID: generic_user_id_attribute_name = os.getenv( "GENERIC_USER_ID_ATTRIBUTE", "preferred_username" @@ -281,16 +334,48 @@ def generic_response_convertor( team_ids = jwt_handler.get_team_ids_from_jwt(cast(dict, response)) all_teams.extend(team_ids) - # Extract user role from SSO response - user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) + # Determine user role based on role_mappings if available + # Only apply role_mappings for GENERIC SSO provider user_role: Optional[LitellmUserRoles] = None - if user_role_from_sso is not None: - role = get_litellm_user_role(user_role_from_sso) - if role is not None: - user_role = role + + if role_mappings is not None and role_mappings.provider.lower() in ["generic", "okta"]: + # Use role_mappings to determine role from groups + group_claim = role_mappings.group_claim + user_groups_raw = get_nested_value(response, group_claim) + + # Handle different formats: could be a list, string (comma-separated), or single value + user_groups: List[str] = [] + if isinstance(user_groups_raw, list): + user_groups = [str(g) for g in user_groups_raw] + elif isinstance(user_groups_raw, str): + # Handle comma-separated string + user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()] + elif user_groups_raw is not None: + # Single value + user_groups = [str(user_groups_raw)] + + if user_groups: + user_role = determine_role_from_groups(user_groups, role_mappings) verbose_proxy_logger.debug( - f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" + f"Determined role '{user_role.value if user_role else None}' from groups '{user_groups}' using role_mappings" ) + else: + # No groups found, use default_role + user_role = role_mappings.default_role + verbose_proxy_logger.debug( + f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}" + ) + + # Fallback to existing logic if role_mappings not used + if user_role is None: + user_role_from_sso = get_nested_value(response, generic_user_role_attribute_name) + if user_role_from_sso is not None: + role = get_litellm_user_role(user_role_from_sso) + if role is not None: + user_role = role + verbose_proxy_logger.debug( + f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'" + ) return CustomOpenID( id=get_nested_value(response, generic_user_id_attribute_name), @@ -369,6 +454,40 @@ async def get_generic_sso_response( userinfo_endpoint=generic_userinfo_endpoint, ) + # Get role_mappings from SSO settings if available + role_mappings: Optional["RoleMappings"] = None + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data + + verbose_proxy_logger.debug( + f"Loaded role_mappings for provider '{role_mappings.provider}'" + ) + except Exception as e: + # If we can't load role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not load role_mappings from database: {e}. Continuing with existing role logic." + ) + def response_convertor(response, client): nonlocal received_response # return for user debugging received_response = response @@ -376,6 +495,7 @@ async def get_generic_sso_response( response=response, jwt_handler=jwt_handler, sso_jwt_handler=sso_jwt_handler, + role_mappings=role_mappings, ) SSOProvider = create_provider( From f4f5ea85dfe7eff7130724b1d7331354468f5ce8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 18 Dec 2025 14:42:41 +0530 Subject: [PATCH 007/330] Add redisvl in requirements.txt --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index c36f94f0752..2fb6c52cfd3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,8 @@ starlette==0.49.1 # starlette fastapi dep backoff==2.2.1 # server dep pyyaml==6.0.2 # server dep uvicorn==0.31.1 # server dep -gunicorn==23.0.0 # server dep +gunicorn==23.0.0 # server depredisvl +redisvl==0.4.1 # redis semantic cache fastuuid==0.13.5 # for uuid4 uvloop==0.21.0 # uvicorn dep, gives us much better performance under load boto3==1.36.0 # aws bedrock/sagemaker calls From 5705aaebbc3247cb1666dec0288ab7b5507b2c79 Mon Sep 17 00:00:00 2001 From: Rens Date: Thu, 18 Dec 2025 15:11:58 +0200 Subject: [PATCH 008/330] Fix Gemini 3 imgs in tool response --- .../prompt_templates/factory.py | 7 +-- ...llm_core_utils_prompt_templates_factory.py | 51 +++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 652692c7b8d..9afea83ef7b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1496,9 +1496,10 @@ def convert_to_gemini_tool_call_result( content_type = content.get("type", "") if content_type == "text": content_str += content.get("text", "") - elif content_type == "input_image": - # Extract image for inline_data (for Computer Use screenshots) - image_url = content.get("image_url", "") + elif content_type in ("input_image", "image_url"): + # Extract image for inline_data (for Computer Use screenshots and tool results) + image_url_data = content.get("image_url", "") + image_url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data if image_url: # Convert image to base64 blob format for Gemini diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 41ac893b4d7..c8fe6efeaa1 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -497,6 +497,57 @@ def test_convert_gemini_messages(): ) +def test_convert_gemini_tool_call_result_with_image_url(): + """ + Test that image_url content type in tool results is handled correctly for Gemini. + Fixes: https://github.com/BerriAI/litellm/issues/18187 + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_result, + ) + from litellm.types.llms.openai import ChatCompletionToolMessage + + # Test with string image_url format + message_str_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_123", + content=[{"type": "image_url", "image_url": "data:image/jpeg;base64,/9j/4AAQ"}], + ) + last_message_with_tool_calls = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "index": 0, + "function": {"name": "get_image", "arguments": "{}"}, + } + ], + } + + result = convert_to_gemini_tool_call_result( + message=message_str_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + # Should have inline_data for the image + assert isinstance(result, list) and any("inline_data" in p for p in result) + + # Test with dict image_url format (OpenAI standard) + message_dict_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_456", + content=[{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ"}}], + ) + last_message_with_tool_calls["tool_calls"][0]["id"] = "call_456" + + result2 = convert_to_gemini_tool_call_result( + message=message_dict_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + assert isinstance(result2, list) and any("inline_data" in p for p in result2) + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly From bf76e66d2ccaace63eb6353ab92e37b57d99b349 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 09:42:59 -0800 Subject: [PATCH 009/330] Working SSO Mapping new user, overrides default user settings --- litellm/proxy/management_endpoints/ui_sso.py | 48 +++++++++++++++++++- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 7a404bcda81..5afccc6fe5f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1173,8 +1173,44 @@ async def insert_sso_user( if user_defined_values is None: raise ValueError("user_defined_values is None") + # Check if role_mappings is configured in SSO settings + role_mappings_configured = False + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + role_mappings_configured = role_mappings_data is not None + except Exception as e: + # If we can't check role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not check role_mappings configuration: {e}. Using default behavior." + ) + + # Apply default_internal_user_params if litellm.default_internal_user_params: - user_defined_values.update(litellm.default_internal_user_params) # type: ignore + # If role_mappings is configured and user_role is already set from SSO, preserve it + if role_mappings_configured and user_defined_values.get("user_role") is not None: + # Preserve the SSO-extracted role, but apply other defaults + preserved_role = user_defined_values.get("user_role") + user_defined_values.update(litellm.default_internal_user_params) # type: ignore + user_defined_values["user_role"] = preserved_role # Restore preserved role + verbose_proxy_logger.debug( + f"Preserved SSO-extracted role '{preserved_role}' (role_mappings configured)" + ) + else: + # Default behavior: update all values including role + user_defined_values.update(litellm.default_internal_user_params) # type: ignore # Set budget for internal users if user_defined_values.get("user_role") == LitellmUserRoles.INTERNAL_USER.value: @@ -1812,7 +1848,15 @@ class SSOAuthenticationHandler: ) user_id = getattr(result, "id", None) user_email = getattr(result, "email", None) - user_role = getattr(result, generic_user_role_attribute_name, None) # type: ignore + if user_role is None: + _role_from_attr = getattr(result, generic_user_role_attribute_name, None) # type: ignore + if _role_from_attr is not None: + # Convert enum to string if needed + user_role = ( + _role_from_attr.value + if isinstance(_role_from_attr, LitellmUserRoles) + else _role_from_attr + ) if user_id is None and result is not None: _first_name = getattr(result, "first_name", "") or "" From 313a613a13e2756308ce828924864815ac36bd38 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 10:56:11 -0800 Subject: [PATCH 010/330] Adding tests --- .../proxy/management_endpoints/test_ui_sso.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 500fc67de89..20829466570 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -3043,3 +3043,108 @@ class TestAddMissingTeamMember: assert set(added_teams) == set( expected_teams_added ), f"Expected teams {expected_teams_added}, but got {added_teams}" + + +@pytest.mark.asyncio +async def test_role_mappings_override_default_internal_user_params(): + """ + Test that when role_mappings is configured in SSO settings, + the SSO-extracted role overrides default_internal_user_params role. + """ + from litellm.proxy._types import NewUserResponse, SSOUserDefinedValues + from litellm.proxy.management_endpoints.ui_sso import insert_sso_user + + # Save original default_internal_user_params + original_default_params = getattr(litellm, "default_internal_user_params", None) + + try: + # Set default_internal_user_params with a role that should be overridden + litellm.default_internal_user_params = { + "user_role": "internal_user", + "max_budget": 100, + "budget_duration": "30d", + "models": ["gpt-3.5-turbo"], + } + + # Mock SSO result + mock_result_openid = CustomOpenID( + id="test-user-123", + email="test@example.com", + display_name="Test User", + provider="microsoft", + team_ids=[], + ) + + # User defined values with SSO-extracted role (from role_mappings) + user_defined_values: SSOUserDefinedValues = { + "user_id": "test-user-123", + "user_email": "test@example.com", + "user_role": "proxy_admin", # Role from SSO role_mappings + "max_budget": None, + "budget_duration": None, + "models": [], + } + + # Mock Prisma client with SSO config that has role_mappings configured + mock_prisma = MagicMock() + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = { + "role_mappings": { + "Admin": "proxy_admin", + "User": "internal_user", + } + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=mock_sso_config + ) + + # Mock new_user function + mock_new_user_response = NewUserResponse( + user_id="test-user-123", + key="sk-xxxxx", + teams=None, + ) + + with patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=mock_prisma, + ), patch( + "litellm.proxy.management_endpoints.ui_sso.new_user", + return_value=mock_new_user_response, + ) as mock_new_user: + # Act + result = await insert_sso_user( + result_openid=mock_result_openid, + user_defined_values=user_defined_values, + ) + + # Assert - verify new_user was called with preserved SSO role + mock_new_user.assert_called_once() + call_args = mock_new_user.call_args + new_user_request = call_args.kwargs["data"] + + # The role from SSO should be preserved, not overridden by default_internal_user_params + assert ( + new_user_request.user_role == "proxy_admin" + ), "SSO-extracted role should override default_internal_user_params role" + + # Other default params should still be applied + assert ( + new_user_request.max_budget == 100 + ), "max_budget from default_internal_user_params should be applied" + assert ( + new_user_request.budget_duration == "30d" + ), "budget_duration from default_internal_user_params should be applied" + + # Note: models are applied via _update_internal_new_user_params inside new_user, + # not in insert_sso_user, so we verify user_defined_values was updated correctly + # by checking that the function completed successfully and other defaults were applied + # The models will be applied when new_user processes the request + + finally: + # Restore original default_internal_user_params + if original_default_params is not None: + litellm.default_internal_user_params = original_default_params + else: + if hasattr(litellm, "default_internal_user_params"): + delattr(litellm, "default_internal_user_params") From ffcac2eebcc3a1f2c74a81171d79c930750f0ee8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 18:04:04 -0800 Subject: [PATCH 011/330] Allow deleting key expiry --- .../key_management_endpoints.py | 13 ++++++-- tests/proxy_unit_tests/test_proxy_utils.py | 5 +++ .../test_key_management_endpoints.py | 31 +++++++++++++++++++ .../KeyLifecycleSettings.tsx | 4 +-- 4 files changed, 48 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 14d221d19e1..cc2ac908149 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -507,7 +507,11 @@ async def _common_key_generation_helper( # noqa: PLR0915 upperbound_duration = duration_in_seconds( duration=upperbound_value ) - user_duration = duration_in_seconds(duration=value) + # Handle special case where duration is "-1" (never expires) + if value == "-1": + user_duration = float('inf') # Infinite duration + else: + user_duration = duration_in_seconds(duration=value) if user_duration > upperbound_duration: raise HTTPException( status_code=400, @@ -1339,7 +1343,10 @@ async def prepare_key_update_data( if "duration" in non_default_values: duration = non_default_values.pop("duration") - if duration and (isinstance(duration, str)) and len(duration) > 0: + if duration == "-1": + # Set expires to None to indicate the key never expires + non_default_values["expires"] = None + elif duration and (isinstance(duration, str)) and len(duration) > 0: duration_s = duration_in_seconds(duration=duration) expires = datetime.now(timezone.utc) + timedelta(seconds=duration_s) non_default_values["expires"] = expires @@ -1452,7 +1459,7 @@ async def update_key_fn( - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - allowed_cache_controls: Optional[list] - List of allowed cache control values - - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) + - duration: Optional[str] - Key validity duration ("30d", "1h", etc.) or "-1" to never expire - permissions: Optional[dict] - Key-specific permissions - send_invite_email: Optional[bool] - Send invite email to user_id - guardrails: Optional[List[str]] - List of active guardrails for the key diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 2e5cfff8bf0..5c3c3948920 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -678,6 +678,11 @@ async def test_prepare_key_update_data(): updated_data = await prepare_key_update_data(data, existing_key_row) assert updated_data["metadata"] is None + # Test duration "-1" sets expires to None (never expires) + data = UpdateKeyRequest(key="test_key", duration="-1") + updated_data = await prepare_key_update_data(data, existing_key_row) + assert updated_data["expires"] is None + @pytest.mark.parametrize( "env_vars, expected_url", diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index ff85e6d9e73..648045a7ea6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -815,6 +815,37 @@ async def test_update_service_account_works_with_team_id(): await prepare_key_update_data(data=data, existing_key_row=existing_key) +@pytest.mark.asyncio +async def test_prepare_key_update_data_duration_never_expires(): + """Test that duration="-1" sets expires to None (never expires).""" + from litellm.proxy._types import UpdateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + prepare_key_update_data, + ) + + # Mock existing key + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=["gpt-3.5-turbo"], + user_id="test-user", + team_id=None, + auto_rotate=False, + rotation_interval=None, + metadata={}, + ) + + # Test setting duration to "-1" (never expires) + update_request = UpdateKeyRequest(key="test-token", duration="-1") + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + # Verify that expires is set to None + assert result["expires"] is None + + @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_requires_team_id(): """ diff --git a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx index 8129f4314fc..81d22b56347 100644 --- a/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx +++ b/ui/litellm-dashboard/src/components/common_components/KeyLifecycleSettings.tsx @@ -64,13 +64,13 @@ const KeyLifecycleSettings: React.FC = ({
Date: Sat, 20 Dec 2025 11:35:20 +0800 Subject: [PATCH 012/330] fix: fixed the issue of handling root paths when processing Discovery protected resource metadata and authorization server metadata URLs. --- .../mcp_server/discoverable_endpoints.py | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffa17a5b7c4..4b6020f582b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -15,6 +15,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.proxy.utils import get_server_root_path router = APIRouter( tags=["mcp"], @@ -381,7 +382,18 @@ async def callback(code: str, state: str): # ------------------------------ # Optional .well-known endpoints for MCP + OAuth discovery # ------------------------------ -@router.get("/.well-known/oauth-protected-resource/{mcp_server_name}/mcp") +""" + Per SEP-985, the client MUST: + 1. Try resource_metadata from WWW-Authenticate header (if present) + 2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path} + ( + If the resource identifier value contains a path or query component, any terminating slash (/) + following the host component MUST be removed before inserting /.well-known/ and the well-known + URI path suffix between the host component and the path(include root path) and/or query components. + https://datatracker.ietf.org/doc/html/rfc9728#section-3.1) + 3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource +""" +@router.get(f"/.well-known/oauth-protected-resource{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}/mcp") @router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp( request: Request, mcp_server_name: Optional[str] = None @@ -403,8 +415,15 @@ async def oauth_protected_resource_mcp( ), # this is what Claude will call } - -@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}") +""" + https://datatracker.ietf.org/doc/html/rfc8414#section-3.1 + RFC 8414: Path-aware OAuth discovery + If the issuer identifier value contains a path component, any + terminating "/" MUST be removed before inserting "/.well-known/" and + the well-known URI suffix between the host component and the path(include root path) + component. +""" +@router.get(f"/.well-known/oauth-authorization-server{'' if get_server_root_path() == '/' else get_server_root_path()}/{{mcp_server_name}}") @router.get("/.well-known/oauth-authorization-server") async def oauth_authorization_server_mcp( request: Request, mcp_server_name: Optional[str] = None From 3a2ab6b0d12be8863d2a7604a1f3e8ab5721b521 Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sat, 20 Dec 2025 11:55:12 +0800 Subject: [PATCH 013/330] fix: added additional grant type into oauth_authorization_server response for fixing mcp auth register bad request issue --- .../proxy/_experimental/mcp_server/discoverable_endpoints.py | 2 +- .../_experimental/mcp_server/test_discoverable_endpoints.py | 3 ++- ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ffa17a5b7c4..5433196dfe3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -428,7 +428,7 @@ async def oauth_authorization_server_mcp( "authorization_endpoint": authorization_endpoint, "token_endpoint": token_endpoint, "response_types_supported": ["code"], - "grant_types_supported": ["authorization_code"], + "grant_types_supported": ["authorization_code", "refresh_token"], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["client_secret_post"], # Claude expects a registration endpoint, even if we just fake it diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6df9abd3fee..30f3d55f028 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -354,7 +354,7 @@ async def test_register_client_remote_registration_success(): request_payload = { "client_name": "Litellm Proxy", - "grant_types": ["authorization_code"], + "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], "token_endpoint_auth_method": "client_secret_post", } @@ -603,6 +603,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): assert response["authorization_endpoint"].startswith("https://litellm.example.com/") assert response["token_endpoint"].startswith("https://litellm.example.com/") assert response["registration_endpoint"].startswith("https://litellm.example.com/") + assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index 9600c962564..a62d8baa75b 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -136,7 +136,7 @@ export const useMcpOAuthFlow = ({ if (!hasPreconfiguredCredentials) { const registration = await registerMcpOAuthClient(accessToken, serverId, { client_name: temporaryPayload.alias || temporaryPayload.server_name || serverId, - grant_types: ["authorization_code"], + grant_types: ["authorization_code", "refresh_token"], response_types: ["code"], token_endpoint_auth_method: temporaryPayload.credentials && temporaryPayload.credentials.client_secret ? "client_secret_post" : "none", From 684fba42eaaf6a4d47795e56fd668b8d46b01525 Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sat, 20 Dec 2025 13:22:29 +0800 Subject: [PATCH 014/330] fix: added RFC RECOMMENDED property(scopes_supported) to protected resource and authorization server metadata --- .../mcp_server/discoverable_endpoints.py | 17 +++++- .../mcp_server/test_discoverable_endpoints.py | 54 ++++++++++++++++++- 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index d6fe3f2b9cf..ded591a8f53 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -398,8 +398,14 @@ async def callback(code: str, state: str): async def oauth_protected_resource_mcp( request: Request, mcp_server_name: Optional[str] = None ): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) # Get the correct base URL considering X-Forwarded-* headers request_base_url = get_request_base_url(request) + mcp_server: Optional[MCPServer] = None + if mcp_server_name: + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name) return { "authorization_servers": [ ( @@ -413,6 +419,7 @@ async def oauth_protected_resource_mcp( if mcp_server_name else f"{request_base_url}/mcp" ), # this is what Claude will call + "scopes_supported": mcp_server.scopes if mcp_server else [], } """ @@ -428,6 +435,9 @@ async def oauth_protected_resource_mcp( async def oauth_authorization_server_mcp( request: Request, mcp_server_name: Optional[str] = None ): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) # Get the correct base URL considering X-Forwarded-* headers request_base_url = get_request_base_url(request) @@ -442,16 +452,21 @@ async def oauth_authorization_server_mcp( else f"{request_base_url}/token" ) + mcp_server: Optional[MCPServer] = None + if mcp_server_name: + mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name) + return { "issuer": request_base_url, # point to your proxy "authorization_endpoint": authorization_endpoint, "token_endpoint": token_endpoint, "response_types_supported": ["code"], + "scopes_supported": mcp_server.scopes if mcp_server else [], "grant_types_supported": ["authorization_code", "refresh_token"], "code_challenge_methods_supported": ["S256"], "token_endpoint_auth_methods_supported": ["client_secret_post"], # Claude expects a registration endpoint, even if we just fake it - "registration_endpoint": f"{request_base_url}/{mcp_server_name}/register", + "registration_endpoint": f"{request_base_url}/{mcp_server_name}/register" if mcp_server_name else f"{request_base_url}/register", } diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 30f3d55f028..4c5723b8284 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -556,9 +556,33 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( oauth_protected_resource_mcp, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) @@ -568,13 +592,14 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): # Call the endpoint response = await oauth_protected_resource_mcp( request=mock_request, - mcp_server_name="test_server", + mcp_server_name="test_oauth", ) # Verify response uses HTTPS URLs assert response["authorization_servers"][0].startswith( "https://litellm.example.com/" ) + assert response["scopes_supported"] == oauth2_server.scopes @pytest.mark.asyncio @@ -584,9 +609,33 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( oauth_authorization_server_mcp, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + from litellm.proxy._types import MCPTransport from fastapi import Request except ImportError: pytest.skip("MCP discoverable endpoints not available") + # Clear registry + global_mcp_server_manager.registry.clear() + + # Create mock OAuth2 server + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="test_client_id", + client_secret="test_client_secret", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + scopes=["read", "write"], + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) @@ -596,7 +645,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): # Call the endpoint response = await oauth_authorization_server_mcp( request=mock_request, - mcp_server_name="test_server", + mcp_server_name="test_oauth", ) # Verify response uses HTTPS URLs @@ -604,6 +653,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): assert response["token_endpoint"].startswith("https://litellm.example.com/") assert response["registration_endpoint"].startswith("https://litellm.example.com/") assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] + assert response["scopes_supported"] == oauth2_server.scopes @pytest.mark.asyncio From 0306f02e74d7fff462fb727639a60e5ae11d64e4 Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sat, 20 Dec 2025 14:13:18 +0800 Subject: [PATCH 015/330] fix: removed initialize the tool name to MCP server name mapping(oauth2) on startup for avoiding 401 error --- .../mcp_server/mcp_server_manager.py | 3 +++ .../mcp_server/test_mcp_server_manager.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8c9d8630457..c2215efe9d0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1913,6 +1913,9 @@ class MCPServerManager: Note: This now handles prefixed tool names """ for server in self.get_registry().values(): + if server.auth_type == MCPAuth.oauth2: + # Skip OAuth2 servers for now as they may require user-specific tokens + continue tools = await self._get_tools_from_server(server) for tool in tools: # The tool.name here is already prefixed from _get_tools_from_server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7a6e5ad17f6..c0ded9c728c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -536,7 +536,26 @@ class TestMCPServerManager: assert ( server.registration_url == "https://discovered.example.com/register" ) + @pytest.mark.asyncio + async def test_config_oauth_initialize_tool_name_to_mcp_server_name_mapping(self): + manager = MCPServerManager() + config = { + "example": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "scopes": ["config"], + "authorization_url": "https://config.example.com/auth", + } + } + + await manager.load_servers_from_config(config) + + # Initialize the tool mapping + await manager._initialize_tool_name_to_mcp_server_name_mapping() + assert manager.tool_name_to_mcp_server_name_mapping == {} + @pytest.mark.asyncio async def test_list_tools_handles_missing_server_alias(self): """Test that list_tools handles servers without alias gracefully""" From 36a369a747fccedb87e12cf26996aebfc221f57e Mon Sep 17 00:00:00 2001 From: Eric84626 Date: Sat, 20 Dec 2025 14:28:52 +0800 Subject: [PATCH 016/330] fix: upgraded mcp sdk depency version for fixing ClosedResourceError --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index f222acc46e6..cb12a658814 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,7 +20,7 @@ google-cloud-aiplatform==1.47.0 # for vertex ai calls google-cloud-iam==2.19.1 # for GCP IAM Redis authentication google-genai==1.22.0 anthropic[vertex]==0.54.0 -mcp==1.21.2 ; python_version >= "3.10" # for MCP server +mcp==1.25.0 ; python_version >= "3.10" # for MCP server google-generativeai==0.5.0 # for vertex ai calls async_generator==1.10.0 # for async ollama calls langfuse==2.59.7 # for langfuse self-hosted logging From b981fddfaadd95d97dca9df4be3561d5a6f9a2cf Mon Sep 17 00:00:00 2001 From: Lucas Rothman Date: Fri, 19 Dec 2025 23:13:36 -0800 Subject: [PATCH 017/330] fix(gemini): properly catch context window exceeded errors Fixes #18282 This PR fixes two issues with Gemini context window error handling: 1. **Pattern matching for Gemini 2.0 Flash**: The previous pattern 'input token count exceeds the maximum number of tokens allowed' doesn't match Gemini 2.0 Flash errors which include dynamic token counts like '(2800010)' in the message. Split into shorter patterns that work with both formats. 2. **Add context window check to Gemini block**: The is_error_str_context_window_exceeded() check was only called for OpenAI-compatible providers, not for Gemini/Vertex AI. Added the check to the Gemini-specific error handling block. Test cases added for both Gemini 2.0 Flash and 2.5/3 error formats. --- .../exception_mapping_utils.py | 12 +++- .../test_exception_mapping_utils.py | 60 ++++++++++++++++++- 2 files changed, 67 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 7bf95ca3404..1517d1e776d 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -78,9 +78,7 @@ class ExceptionCheckers: "is longer than the model's context length", "input tokens exceed the configured limit", "`inputs` tokens + `max_new_tokens` must be", - # Gemini pattern: "The input token count exceeds the maximum number of tokens allowed" - # See: https://github.com/BerriAI/litellm/issues/XXXX - "input token count exceeds the maximum number of tokens allowed", + "exceeds the maximum number of tokens allowed", # Gemini ] for substring in known_exception_substrings: if substring in _error_str_lowercase: @@ -1262,6 +1260,14 @@ def exception_type( # type: ignore # noqa: PLR0915 model=model, llm_provider=custom_llm_provider, ) + elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): + exception_mapping_worked = True + raise ContextWindowExceededError( + message=f"ContextWindowExceededError: {custom_llm_provider.capitalize()}Exception - {error_str}", + model=model, + llm_provider=custom_llm_provider, + litellm_debug_info=extra_information, + ) elif ( "None Unknown Error." in error_str or "Content has no parts." in error_str diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index f69b9c35236..9e742a83c6a 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -40,8 +40,7 @@ context_window_test_cases = [ "`inputs` tokens + `max_new_tokens` must be <= 4096", True, ), - # Gemini context window error format - # See: https://github.com/BerriAI/litellm/issues/XXXX + # Gemini 2.5/3 format ( "The input token count exceeds the maximum number of tokens allowed 1048576.", True, @@ -50,6 +49,15 @@ context_window_test_cases = [ "GeminiException BadRequestError - {\n \"error\": {\n \"code\": 400,\n \"message\": \"The input token count exceeds the maximum number of tokens allowed 1048576.\",\n \"status\": \"INVALID_ARGUMENT\"\n }\n}\n", True, ), + # Gemini 2.0 Flash format (includes input token count in message) + ( + "The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).", + True, + ), + ( + "GeminiException BadRequestError - {\n \"error\": {\n \"code\": 400,\n \"message\": \"The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).\",\n \"status\": \"INVALID_ARGUMENT\"\n }\n}\n", + True, + ), # Test case insensitivity ("ERROR: THIS MODEL'S MAXIMUM CONTEXT LENGTH IS 1024.", True), # Cerebras context window error format @@ -169,6 +177,54 @@ class TestExceptionCheckers: result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is False, f"Should NOT detect policy violation in: {error_str}" +gemini_context_window_test_cases = [ + # Gemini 2.0 Flash format (includes input token count in message) + ( + "The input token count (2800010) exceeds the maximum number of tokens allowed (1048575).", + True, + ), + # Gemini 2.5/3 format + ( + "The input token count exceeds the maximum number of tokens allowed (1048576).", + True, + ), + ("A generic error occurred.", False), +] + + +@pytest.mark.parametrize( + "error_message, should_raise_context_window", gemini_context_window_test_cases +) +def test_gemini_context_window_error_mapping(error_message, should_raise_context_window): + """ + Tests that the exception_type function correctly maps Gemini's + context window exceeded errors to litellm.ContextWindowExceededError. + """ + model = "gemini/gemini-2.0-flash" + custom_llm_provider = "gemini" + + # Create a generic exception with the specific error message + original_exception = Exception(error_message) + + if should_raise_context_window: + with pytest.raises(litellm.ContextWindowExceededError) as excinfo: + exception_type( + model=model, + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + ) + # Check if the raised exception is indeed a ContextWindowExceededError + assert isinstance(excinfo.value, litellm.ContextWindowExceededError) + else: + # For the negative case, we expect it to raise a generic APIConnectionError + with pytest.raises(litellm.APIConnectionError): + exception_type( + model=model, + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + ) + + # Test cases for Vertex AI RateLimitError mapping # As per https://github.com/BerriAI/litellm/issues/16189 vertex_rate_limit_test_cases = [ From 9002f75277228c0723d0d503f71f1f1a5de4638f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 20 Dec 2025 11:01:41 -0600 Subject: [PATCH 018/330] Require auth for MCP connection test --- .../mcp_server/rest_endpoints.py | 8 +-- .../mcp_server/test_rest_endpoints.py | 55 +++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 032331ece02..92c390f0a64 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,4 @@ import importlib -import traceback from typing import Dict, List, Optional, Union from fastapi import APIRouter, Depends, Query, Request @@ -329,16 +328,15 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) - stack_trace = traceback.format_exc() return { "status": "error", - "message": f"An internal error has occurred: {str(e)}", - "stack_trace": stack_trace, + "message": "An internal error has occurred while testing the MCP server.", } - @router.post("/test/connection") + @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) async def test_connection( request: NewMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Test if we can connect to the provided MCP server before adding it diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index a0c09663a88..85ec807b1ff 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -7,6 +7,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth from litellm.types.mcp import MCPAuth @@ -31,6 +32,60 @@ def _build_request(headers: Optional[Dict[str, str]] = None) -> Request: return Request(scope, receive=receive) +def _get_route(path: str, method: str): + for route in rest_endpoints.router.routes: + if getattr(route, "path", None) == path and method in getattr( + route, "methods", set() + ): + return route + raise AssertionError(f"Route {method} {path} not found") + + +def _route_has_dependency(route, dependency) -> bool: + if any( + getattr(dep, "dependency", None) == dependency + for dep in getattr(route, "dependencies", []) + ): + return True + dependant = getattr(route, "dependant", None) + if dependant is None: + return False + return any(getattr(dep, "call", None) == dependency for dep in dependant.dependencies) + + +@pytest.mark.asyncio +async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): + def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def failing_operation(client): + raise RuntimeError("boom") + + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, failing_operation + ) + + assert result["status"] == "error" + assert "stack_trace" not in result + + +def test_test_connection_requires_auth_dependency(): + route = _get_route("/mcp-rest/test/connection", "POST") + assert _route_has_dependency(route, user_api_key_auth) + + @pytest.mark.asyncio async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch): """Ensure credential-based auth forwards the auth_value to the MCP client.""" From acce6b9c83f143116038b49be710398802cd540e Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sat, 20 Dec 2025 15:56:46 -0600 Subject: [PATCH 019/330] Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/test_rest_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 85ec807b1ff..31ab4afb631 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -82,7 +82,7 @@ async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): def test_test_connection_requires_auth_dependency(): - route = _get_route("/mcp-rest/test/connection", "POST") + route = _get_route("/test/connection", "POST") assert _route_has_dependency(route, user_api_key_auth) From b2588dc39910a50cf6b348a974d2f53fabd3bc1b Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 21 Dec 2025 22:40:27 +0530 Subject: [PATCH 020/330] fix: lost tool_calls when streaming has both text and tool_calls --- .../adapters/transformation.py | 8 +-- ...al_pass_through_adapters_transformation.py | 53 +++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 9cfbf1b6d8d..8868fabdcef 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -740,9 +740,7 @@ class LiteLLMAnthropicMessagesAdapter: from litellm.types.llms.anthropic import TextBlock, ToolUseBlock for choice in choices: - if choice.delta.content is not None and len(choice.delta.content) > 0: - return "text", TextBlock(type="text", text="") - elif ( + if ( choice.delta.tool_calls is not None and len(choice.delta.tool_calls) > 0 and choice.delta.tool_calls[0].function is not None @@ -753,6 +751,8 @@ class LiteLLMAnthropicMessagesAdapter: name=choice.delta.tool_calls[0].function.name or "", input={}, # type: ignore[typeddict-item] ) + elif choice.delta.content is not None and len(choice.delta.content) > 0: + return "text", TextBlock(type="text", text="") elif isinstance(choice, StreamingChoices) and hasattr( choice.delta, "thinking_blocks" ): @@ -796,7 +796,7 @@ class LiteLLMAnthropicMessagesAdapter: for choice in choices: if choice.delta.content is not None and len(choice.delta.content) > 0: text += choice.delta.content - elif choice.delta.tool_calls is not None: + if choice.delta.tool_calls is not None: partial_json = "" for tool in choice.delta.tool_calls: if ( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 9d6fbf66e48..6aadbc058d1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1055,3 +1055,56 @@ def test_translate_anthropic_messages_to_openai_tool_result_single_item_backward f"got {type(tool_message['content'])}" ) assert tool_message["content"] == "72°F and sunny" + + +def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238(): + """ + When a streaming choice contains both text content and tool_calls, + both should be processed (tool_calls should not be ignored). + """ + # streaming choice with both text and tool_calls + choices = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content="Here is some text for litellm", + role=None, + function_call=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="toolu_bdrk_013xRVejhv3ybmLEGCoZib2b", + function=Function(arguments='{"cmd": "init"}', name="Bash"), + type="function", + index=0, + ) + ], + audio=None, + ), + logprobs=None, + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + + # When both text and tool_calls exist, tool_calls (input_json_delta) takes priority + ( + type_of_content, + content_block_delta, + ) = adapter._translate_streaming_openai_chunk_to_anthropic(choices=choices) + + assert type_of_content == "input_json_delta" + assert content_block_delta["partial_json"] == '{"cmd": "init"}' + + # When both text and tool_calls exist, tool_use should be detected and tool name captured + ( + block_type, + content_block_start, + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "tool_use" + assert content_block_start["name"] == "Bash" + assert content_block_start["id"] == "toolu_bdrk_013xRVejhv3ybmLEGCoZib2b" From 34b500c7f5b24b6dd24c9c286081aa07c56b6d4b Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 22 Dec 2025 11:21:59 +0900 Subject: [PATCH 021/330] feat: support MCP stdio header env overrides --- docs/my-website/docs/mcp.md | 27 +++++- .../mcp_server/mcp_server_manager.py | 62 ++++++++++++- .../mcp_server/rest_endpoints.py | 58 ++++++++---- .../proxy/_experimental/mcp_server/server.py | 6 ++ tests/mcp_tests/test_mcp_server.py | 16 +++- .../mcp_server/test_mcp_server.py | 58 ++++++++++-- .../mcp_server/test_mcp_server_manager.py | 93 +++++++++++++++++-- .../mcp_server/test_rest_endpoints.py | 16 +++- 8 files changed, 295 insertions(+), 41 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index f9c9cbb4562..a70e3d24188 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -746,8 +746,33 @@ curl --location 'http://localhost:4000/github_mcp/mcp' \ 3. **Header Forwarding**: LiteLLM automatically forwards matching headers to the backend MCP server 4. **Authentication**: The backend MCP server receives both the configured auth headers and the custom headers ---- +### Passing Request Headers to STDIO env Vars + +If your stdio MCP server needs per-request credentials, you can map HTTP headers from the client request directly into the environment for the launched stdio process. Reference the header name in the env value using the `${X-HEADER_NAME}` syntax. LiteLLM will read that header from the incoming request and set the env var before starting the command. + +```json title="Forward X-GITHUB_PERSONAL_ACCESS_TOKEN header to stdio env" showLineNumbers +{ + "mcpServers": { + "github": { + "command": "docker", + "args": [ + "run", + "-i", + "--rm", + "-e", + "GITHUB_PERSONAL_ACCESS_TOKEN", + "ghcr.io/github/github-mcp-server" + ], + "env": { + "GITHUB_PERSONAL_ACCESS_TOKEN": "${X-GITHUB_PERSONAL_ACCESS_TOKEN}" + } + } + } +} +``` + +In this example, when a client makes a request with the `X-GITHUB_PERSONAL_ACCESS_TOKEN` header, the proxy forwards that value into the stdio process as the `GITHUB_PERSONAL_ACCESS_TOKEN` environment variable. ## Using your MCP with client side credentials diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 8c9d8630457..2260649e8b2 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -84,6 +84,8 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: class MCPServerManager: + _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") + def __init__(self): self.registry: Dict[str, MCPServer] = {} self.config_mcp_servers: Dict[str, MCPServer] = {} @@ -671,11 +673,39 @@ class MCPServerManager: ######################################################### # Methods that call the upstream MCP servers ######################################################### + def _build_stdio_env( + self, + server: MCPServer, + raw_headers: Optional[Dict[str, str]] = None, + ) -> Optional[Dict[str, str]]: + """Resolve stdio env values, supporting header-driven placeholders.""" + + if server.transport != MCPTransport.stdio or not server.env: + return None + + resolved_env: Dict[str, str] = {} + normalized_headers = {k.lower(): v for k, v in (raw_headers or {}).items()} + + for env_key, env_value in server.env.items(): + stripped_value = env_value.strip() + match = self._STDIO_ENV_TEMPLATE_PATTERN.match(stripped_value) + if match: + header_name = match.group(1) + header_value = normalized_headers.get(header_name.lower()) + if header_value is None: + continue + resolved_env[env_key] = header_value + else: + resolved_env[env_key] = env_value + + return resolved_env + def _create_mcp_client( self, server: MCPServer, mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, + stdio_env: Optional[Dict[str, str]] = None, ) -> MCPClient: """ Create an MCPClient instance for the given server. @@ -692,10 +722,13 @@ class MCPServerManager: # Handle stdio transport if transport == MCPTransport.stdio: # For stdio, we need to get the stdio config from the server + resolved_env = stdio_env if stdio_env is not None else server.env or {} stdio_config: Optional[MCPStdioConfig] = None if server.command and server.args is not None: stdio_config = MCPStdioConfig( - command=server.command, args=server.args, env=server.env or {} + command=server.command, + args=server.args, + env=resolved_env, ) return MCPClient( @@ -725,6 +758,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -751,10 +785,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) ## HANDLE OPENAPI TOOLS @@ -784,6 +821,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[Prompt]: """ Helper method to get prompts from a single MCP server with prefixed names. @@ -807,10 +845,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) prompts = await client.list_prompts() @@ -833,6 +874,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[Resource]: """Fetch available resources from a single MCP server.""" @@ -847,10 +889,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) resources = await client.list_resources() @@ -873,6 +918,7 @@ class MCPServerManager: mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, + raw_headers: Optional[Dict[str, str]] = None, ) -> List[ResourceTemplate]: """Fetch available resource templates from a single MCP server.""" @@ -887,10 +933,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) resource_templates = await client.list_resource_templates() @@ -913,6 +962,7 @@ class MCPServerManager: url: AnyUrl, mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> ReadResourceResult: """Read resource contents from a specific MCP server.""" @@ -924,10 +974,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) return await client.read_resource(url) @@ -939,6 +992,7 @@ class MCPServerManager: arguments: Optional[Dict[str, Any]] = None, mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, extra_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ) -> GetPromptResult: """Fetch a specific prompt definition from a single MCP server.""" @@ -950,10 +1004,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + stdio_env = self._build_stdio_env(server, raw_headers) + client = self._create_mcp_client( server=server, mcp_auth_header=mcp_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) get_prompt_request_params = GetPromptRequestParams( @@ -1742,10 +1799,13 @@ class MCPServerManager: extra_headers = {} extra_headers.update(mcp_server.static_headers) + stdio_env = self._build_stdio_env(mcp_server, raw_headers) + client = self._create_mcp_client( server=mcp_server, mcp_auth_header=server_auth_header, extra_headers=extra_headers, + stdio_env=stdio_env, ) call_tool_params = MCPCallToolRequestParams( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 032331ece02..891b52db7af 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -71,12 +71,17 @@ if MCP_AVAILABLE: for tool in tools ] - async def _get_tools_for_single_server(server, server_auth_header): + async def _get_tools_for_single_server( + server, + server_auth_header, + raw_headers: Optional[Dict[str, str]] = None, + ): """Helper function to get tools for a single server.""" tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, add_prefix=False, + raw_headers=raw_headers, ) # Filter tools based on allowed_tools configuration @@ -122,6 +127,7 @@ if MCP_AVAILABLE: try: # Extract auth headers from request headers = request.headers + raw_headers_from_request = dict(headers) mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( headers ) @@ -148,7 +154,7 @@ if MCP_AVAILABLE: try: list_tools_result = await _get_tools_for_single_server( - server, server_auth_header + server, server_auth_header, raw_headers_from_request ) except Exception as e: verbose_logger.exception( @@ -169,7 +175,7 @@ if MCP_AVAILABLE: try: tools_result = await _get_tools_for_single_server( - server, server_auth_header + server, server_auth_header, raw_headers_from_request ) list_tools_result.extend(tools_result) except Exception as e: @@ -232,13 +238,13 @@ if MCP_AVAILABLE: # but they weren't being extracted and passed to call_mcp_tool. # This fix ensures auth headers are properly extracted from the HTTP request # and passed through to the MCP server for authentication. + headers = request.headers + raw_headers_from_request = dict(headers) mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( - request.headers + headers ) mcp_server_auth_headers = ( - MCPRequestHandler._get_mcp_server_auth_headers_from_headers( - request.headers - ) + MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers) ) # Add extracted headers to data dict to pass to call_mcp_tool @@ -246,6 +252,7 @@ if MCP_AVAILABLE: data["mcp_auth_header"] = mcp_auth_header if mcp_server_auth_headers: data["mcp_server_auth_headers"] = mcp_server_auth_headers + data["raw_headers"] = raw_headers_from_request result = await call_mcp_tool(**data) return result @@ -300,6 +307,7 @@ if MCP_AVAILABLE: operation, mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, oauth2_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, ): """ Common helper to create MCP client, execute operation, and ensure proper cleanup. @@ -312,17 +320,27 @@ if MCP_AVAILABLE: Operation result or error response """ try: + server_model = MCPServer( + server_id=request.server_id or "", + name=request.alias or request.server_name or "", + url=request.url, + transport=request.transport, + auth_type=request.auth_type, + mcp_info=request.mcp_info, + command=request.command, + args=request.args, + env=request.env, + ) + + stdio_env = global_mcp_server_manager._build_stdio_env( + server_model, raw_headers + ) + client = global_mcp_server_manager._create_mcp_client( - server=MCPServer( - server_id=request.server_id or "", - name=request.alias or request.server_name or "", - url=request.url, - transport=request.transport, - auth_type=request.auth_type, - mcp_info=request.mcp_info, - ), + server=server_model, mcp_auth_header=mcp_auth_header, extra_headers=oauth2_headers, + stdio_env=stdio_env, ) return await operation(client) @@ -338,7 +356,8 @@ if MCP_AVAILABLE: @router.post("/test/connection") async def test_connection( - request: NewMCPServerRequest, + request: Request, + new_mcp_server_request: NewMCPServerRequest, ): """ Test if we can connect to the provided MCP server before adding it @@ -351,7 +370,11 @@ if MCP_AVAILABLE: await client.run_with_session(_noop) return {"status": "ok"} - return await _execute_with_mcp_client(request, _test_connection_operation) + return await _execute_with_mcp_client( + new_mcp_server_request, + _test_connection_operation, + raw_headers=dict(request.headers), + ) @router.post("/test/tools/list") async def test_tools_list( @@ -405,4 +428,5 @@ if MCP_AVAILABLE: _list_tools_operation, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, + raw_headers=dict(request.headers), ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index bdff60c932b..e00fdbfb930 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -775,6 +775,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=add_prefix, + raw_headers=raw_headers, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) @@ -854,6 +855,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=add_prefix, + raw_headers=raw_headers, ) all_prompts.extend(prompts) @@ -912,6 +914,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=add_prefix, + raw_headers=raw_headers, ) all_resources.extend(resources) @@ -969,6 +972,7 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, extra_headers=extra_headers, add_prefix=add_prefix, + raw_headers=raw_headers, ) ) all_resource_templates.extend(resource_templates) @@ -1392,6 +1396,7 @@ if MCP_AVAILABLE: arguments=arguments, mcp_auth_header=server_auth_header, extra_headers=extra_headers, + raw_headers=raw_headers, ) async def mcp_read_resource( @@ -1440,6 +1445,7 @@ if MCP_AVAILABLE: url=url, mcp_auth_header=server_auth_header, extra_headers=extra_headers, + raw_headers=raw_headers, ) def _get_standard_logging_mcp_tool_call( diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index d3112714a9c..9242dfc75f4 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -812,10 +812,19 @@ async def test_get_tools_from_mcp_servers(): return_value=["server1_id", "server2_id"] ) mock_manager_2.get_mcp_server_by_id = lambda server_id: mock_server_1 if server_id == "server1_id" else mock_server_2 + async def mock_get_tools_side_effect( + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=False, + raw_headers=None, + ): + if server.server_id == "server1_id": + return [mock_tool_1] + return [mock_tool_2] + mock_manager_2._get_tools_from_server = AsyncMock( - side_effect=lambda server, mcp_auth_header=None, extra_headers=None, add_prefix=False: ( - [mock_tool_1] if server.server_id == "server1_id" else [mock_tool_2] - ) + side_effect=mock_get_tools_side_effect ) with patch( @@ -1693,6 +1702,7 @@ async def test_get_tools_for_single_server(): server=mock_server, mcp_auth_header="Bearer test_token", add_prefix=False, + raw_headers=None, ) # Verify the result diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 4fc94000d61..a1fbddec586 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -294,6 +294,7 @@ async def test_mcp_get_prompt_success(): arguments={"foo": "bar"}, mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, + raw_headers=None, ) assert result is prompt_result @@ -349,6 +350,7 @@ async def test_mcp_read_resource_success(): url="https://example.com/resource", mcp_auth_header={"Authorization": "token"}, extra_headers={"X-Test": "1"}, + raw_headers=None, ) assert result is read_result @@ -428,7 +430,11 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): ) async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=True + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=True, + raw_headers=None, ): if server.name == "working_server": # Working server returns tools @@ -524,7 +530,11 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): ) async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=True + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=True, + raw_headers=None, ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -839,13 +849,19 @@ async def test_oauth2_headers_passed_to_mcp_client(): # This will capture the arguments passed to _create_mcp_client captured_client_args = {} - def mock_create_mcp_client(server, mcp_auth_header=None, extra_headers=None): + def mock_create_mcp_client( + server, + mcp_auth_header=None, + extra_headers=None, + stdio_env=None, + ): # Capture the arguments for verification captured_client_args.update( { "server": server, "mcp_auth_header": mcp_auth_header, "extra_headers": extra_headers, + "stdio_env": stdio_env, } ) # Return a mock client that doesn't actually connect @@ -934,7 +950,11 @@ async def test_list_tools_single_server_unprefixed_names(): mock_manager.get_mcp_server_by_id = MagicMock(return_value=server) async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=False + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=False, + raw_headers=None, ): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" @@ -1006,7 +1026,11 @@ async def test_list_tools_multiple_servers_prefixed_names(): ) async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=True + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=True, + raw_headers=None, ): tool = MagicMock() # When multiple servers, add_prefix should be True -> prefixed names @@ -1147,7 +1171,11 @@ async def test_list_tools_filters_by_key_team_permissions(): mock_manager.get_mcp_server_by_id = lambda server_id: server async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=False + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=False, + raw_headers=None, ): # Return 4 tools, but only 2 should be allowed tool1 = MagicMock() @@ -1248,7 +1276,11 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): mock_manager.get_mcp_server_by_id = lambda server_id: server async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=False + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=False, + raw_headers=None, ): # Return 4 tools tool1 = MagicMock() @@ -1334,7 +1366,11 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): mock_manager.get_mcp_server_by_id = lambda server_id: server async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=False + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=False, + raw_headers=None, ): # Return 3 tools tool1 = MagicMock() @@ -1423,7 +1459,11 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): mock_manager.get_mcp_server_by_id = MagicMock(return_value=server) async def mock_get_tools_from_server( - server, mcp_auth_header=None, extra_headers=None, add_prefix=True + server, + mcp_auth_header=None, + extra_headers=None, + add_prefix=True, + raw_headers=None, ): # Return tools WITH prefix (as they come from MCP server) tool1 = MagicMock() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7a6e5ad17f6..ff016a1a130 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -8,6 +8,7 @@ from fastapi import HTTPException # Add the parent directory to the path so we can import litellm sys.path.insert(0, "../../../../../") + import httpx from mcp import ReadResourceResult, Resource from mcp.types import ( @@ -99,6 +100,53 @@ class TestMCPServerManager: assert client.stdio_config["args"] == ["server.js"] assert client.stdio_config["env"] == {"NODE_ENV": "test"} + def test_build_stdio_env_only_accepts_x_prefixed_placeholders(self): + """Ensure only ${X-*} placeholders are substituted from headers.""" + manager = MCPServerManager() + server = MCPServer( + server_id="stdio-server-env", + name="stdio_env", + transport=MCPTransport.stdio, + command="node", + args=["server.js"], + env={ + "PASSTHROUGH": "${X-Test-Header}", + "STATIC": "value", + "IGNORED": "${Not-Allowed}", + }, + ) + + env = manager._build_stdio_env( + server, + raw_headers={ + "x-test-header": "resolved-value", + "x-not-used": "other", + }, + ) + + assert env == { + "PASSTHROUGH": "resolved-value", + "STATIC": "value", + "IGNORED": "${Not-Allowed}", + } + + def test_build_stdio_env_missing_header_skips_entry(self): + """Ensure missing headers drop the placeholder from the resolved env.""" + manager = MCPServerManager() + server = MCPServer( + server_id="stdio-server-env-miss", + name="stdio_env_miss", + transport=MCPTransport.stdio, + command="node", + args=["server.js"], + env={"EXPECTED": "${X-Missing}"}, + ) + + env = manager._build_stdio_env(server, raw_headers={}) + + # When the header isn't provided, the key is omitted entirely + assert env == {} + @pytest.mark.asyncio async def test_list_tools_with_server_specific_auth_headers(self): """Test list_tools method with server-specific auth headers""" @@ -123,7 +171,10 @@ class TestMCPServerManager: # Mock _get_tools_from_server to return different results async def mock_get_tools_from_server( - server, mcp_auth_header=None, mcp_protocol_version=None + server, + mcp_auth_header=None, + mcp_protocol_version=None, + raw_headers=None, ): if server.name == "github": tool1 = MagicMock() @@ -174,7 +225,10 @@ class TestMCPServerManager: # Mock _get_tools_from_server async def mock_get_tools_from_server( - server, mcp_auth_header=None, mcp_protocol_version=None + server, + mcp_auth_header=None, + mcp_protocol_version=None, + raw_headers=None, ): assert mcp_auth_header == "legacy-token" # Should use legacy header tool = MagicMock() @@ -209,7 +263,10 @@ class TestMCPServerManager: # Mock _get_tools_from_server async def mock_get_tools_from_server( - server, mcp_auth_header=None, mcp_protocol_version=None + server, + mcp_auth_header=None, + mcp_protocol_version=None, + raw_headers=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -373,6 +430,7 @@ class TestMCPServerManager: server=server, mcp_auth_header="auth", extra_headers=None, + stdio_env=None, ) mock_client.list_resource_templates.assert_awaited_once() mock_prefix.assert_called_once_with(mock_templates, server, add_prefix=False) @@ -554,7 +612,10 @@ class TestMCPServerManager: # Mock _get_tools_from_server async def mock_get_tools_from_server( - server, mcp_auth_header=None, mcp_protocol_version=None + server, + mcp_auth_header=None, + mcp_protocol_version=None, + raw_headers=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -587,7 +648,11 @@ class TestMCPServerManager: manager.get_mcp_server_by_id = MagicMock(return_value=server) # Mock successful _get_tools_from_server - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, + mcp_auth_header=None, + raw_headers=None, + ): tool1 = MagicMock() tool1.name = "tool1" tool2 = MagicMock() @@ -621,7 +686,11 @@ class TestMCPServerManager: manager.get_mcp_server_by_id = MagicMock(return_value=server) # Mock failed _get_tools_from_server - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, + mcp_auth_header=None, + raw_headers=None, + ): raise Exception("Connection timeout") manager._get_tools_from_server = mock_get_tools_from_server @@ -683,7 +752,11 @@ class TestMCPServerManager: manager.get_mcp_server_by_id = mock_get_server_by_id # Mock _get_tools_from_server with different results - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, + mcp_auth_header=None, + raw_headers=None, + ): if server.server_id == "server1": tool = MagicMock() tool.name = "tool1" @@ -724,7 +797,11 @@ class TestMCPServerManager: manager.get_mcp_server_by_id = MagicMock(return_value=server) # Mock _get_tools_from_server to verify auth header is passed - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, + mcp_auth_header=None, + raw_headers=None, + ): assert mcp_auth_header == "test-token" tool = MagicMock() tool.name = "tool1" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index a0c09663a88..ce38ee59e4d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -37,7 +37,13 @@ async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch): captured: dict = {} - async def fake_execute(request, operation, mcp_auth_header=None, oauth2_headers=None): + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): captured["mcp_auth_header"] = mcp_auth_header captured["oauth2_headers"] = oauth2_headers return { @@ -87,7 +93,13 @@ async def test_test_tools_list_extracts_oauth2_headers(monkeypatch): captured: dict = {} - async def fake_execute(request, operation, mcp_auth_header=None, oauth2_headers=None): + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): captured["mcp_auth_header"] = mcp_auth_header captured["oauth2_headers"] = oauth2_headers return { From 2afa5fc9fb72984d34f8ed32c296fc8b00303d48 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 11:03:19 -0800 Subject: [PATCH 022/330] ruff check and mypy linting --- litellm/proxy/management_endpoints/ui_sso.py | 126 ++++++++++++------- 1 file changed, 80 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 013859e1d56..dc976e1ce64 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -391,20 +391,8 @@ def generic_response_convertor( ) -async def get_generic_sso_response( - request: Request, - jwt_handler: JWTHandler, - sso_jwt_handler: Optional[ - JWTHandler - ], # sso specific jwt handler - used for restricted sso group access control - generic_client_id: str, - redirect_url: str, -) -> Tuple[Union[OpenID, dict], Optional[dict]]: # return received response - # make generic sso provider - from fastapi_sso.sso.base import DiscoveryDocument - from fastapi_sso.sso.generic import create_provider - - received_response: Optional[dict] = None +def _setup_generic_sso_env_vars(generic_client_id: str, redirect_url: str) -> Tuple[str, List[str], str, str, str, bool]: + """Setup and validate Generic SSO environment variables.""" generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None) generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ") generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None) @@ -413,6 +401,8 @@ async def get_generic_sso_response( generic_include_client_id = ( os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true" ) + + # Validate required environment variables if generic_client_secret is None: raise ProxyException( message="GENERIC_CLIENT_SECRET not set. Set it in .env file", @@ -441,6 +431,7 @@ async def get_generic_sso_response( param="GENERIC_USERINFO_ENDPOINT", code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + verbose_proxy_logger.debug( f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}" ) @@ -448,6 +439,80 @@ async def get_generic_sso_response( f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n" ) + return ( + generic_client_secret, + generic_scope, + generic_authorization_endpoint, + generic_token_endpoint, + generic_userinfo_endpoint, + generic_include_client_id, + ) + + +async def _setup_role_mappings() -> Optional["RoleMappings"]: + """Setup role mappings from SSO database settings.""" + role_mappings: Optional["RoleMappings"] = None + try: + from litellm.proxy.utils import get_prisma_client_or_throw + + prisma_client = get_prisma_client_or_throw( + "Prisma client is None, connect a database to your proxy" + ) + + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + + if sso_db_record and sso_db_record.sso_settings: + sso_settings_dict = dict(sso_db_record.sso_settings) + role_mappings_data = sso_settings_dict.get("role_mappings") + + if role_mappings_data: + from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings + if isinstance(role_mappings_data, dict): + role_mappings = RoleMappings(**role_mappings_data) + elif isinstance(role_mappings_data, RoleMappings): + role_mappings = role_mappings_data + + if role_mappings: + verbose_proxy_logger.debug( + f"Loaded role_mappings for provider '{role_mappings.provider}'" + ) + except Exception as e: + # If we can't load role_mappings, continue with existing logic + verbose_proxy_logger.debug( + f"Could not load role_mappings from database: {e}. Continuing with existing role logic." + ) + + return role_mappings + + +async def get_generic_sso_response( + request: Request, + jwt_handler: JWTHandler, + sso_jwt_handler: Optional[ + JWTHandler + ], # sso specific jwt handler - used for restricted sso group access control + generic_client_id: str, + redirect_url: str, +) -> Tuple[Union[OpenID, dict], Optional[dict]]: # return received response + # make generic sso provider + from fastapi_sso.sso.base import DiscoveryDocument + from fastapi_sso.sso.generic import create_provider + + received_response: Optional[dict] = None + + # Setup environment variables + ( + generic_client_secret, + generic_scope, + generic_authorization_endpoint, + generic_token_endpoint, + generic_userinfo_endpoint, + generic_include_client_id, + ) = _setup_generic_sso_env_vars(generic_client_id, redirect_url) + discovery = DiscoveryDocument( authorization_endpoint=generic_authorization_endpoint, token_endpoint=generic_token_endpoint, @@ -455,38 +520,7 @@ async def get_generic_sso_response( ) # Get role_mappings from SSO settings if available - role_mappings: Optional["RoleMappings"] = None - try: - from litellm.proxy.utils import get_prisma_client_or_throw - - prisma_client = get_prisma_client_or_throw( - "Prisma client is None, connect a database to your proxy" - ) - - # Get SSO config from dedicated table - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( - where={"id": "sso_config"} - ) - - if sso_db_record and sso_db_record.sso_settings: - sso_settings_dict = dict(sso_db_record.sso_settings) - role_mappings_data = sso_settings_dict.get("role_mappings") - - if role_mappings_data: - from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings - if isinstance(role_mappings_data, dict): - role_mappings = RoleMappings(**role_mappings_data) - elif isinstance(role_mappings_data, RoleMappings): - role_mappings = role_mappings_data - - verbose_proxy_logger.debug( - f"Loaded role_mappings for provider '{role_mappings.provider}'" - ) - except Exception as e: - # If we can't load role_mappings, continue with existing logic - verbose_proxy_logger.debug( - f"Could not load role_mappings from database: {e}. Continuing with existing role logic." - ) + role_mappings = await _setup_role_mappings() def response_convertor(response, client): nonlocal received_response # return for user debugging From 30fa90f70da207cd576861e4e8a705c285a68c61 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Mon, 22 Dec 2025 11:24:30 -0800 Subject: [PATCH 023/330] [Feat] Enable async_post_call_failure_hook to transform error responses (#18348) --- docs/my-website/docs/proxy/call_hooks.md | 52 ++++++- litellm/integrations/custom_logger.py | 17 +- litellm/proxy/auth/auth_exception_handler.py | 19 +-- litellm/proxy/common_request_processing.py | 12 +- litellm/proxy/utils.py | 32 +++- .../proxy/auth/test_user_api_key_auth.py | 4 +- ...test_post_call_failure_hook_integration.py | 146 ++++++++++++++++++ .../test_llm_pass_through_endpoints.py | 4 +- 8 files changed, 263 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_post_call_failure_hook_integration.py diff --git a/docs/my-website/docs/proxy/call_hooks.md b/docs/my-website/docs/proxy/call_hooks.md index fa420009cf1..fe865f67e09 100644 --- a/docs/my-website/docs/proxy/call_hooks.md +++ b/docs/my-website/docs/proxy/call_hooks.md @@ -17,6 +17,7 @@ import Image from '@theme/IdealImage'; | `async_pre_call_hook` | Modify incoming request before it's sent to model | Before the LLM API call is made | | `async_moderation_hook` | Run checks on input in parallel to LLM API call | In parallel with the LLM API call | | `async_post_call_success_hook` | Modify outgoing response (non-streaming) | After successful LLM API call, for non-streaming responses | +| `async_post_call_failure_hook` | Transform error responses sent to clients | After failed LLM API call | | `async_post_call_streaming_hook` | Modify outgoing response (streaming) | After successful LLM API call, for streaming responses | See a complete example with our [parallel request rate limiter](https://github.com/BerriAI/litellm/blob/main/litellm/proxy/hooks/parallel_request_limiter.py) @@ -60,7 +61,21 @@ class MyCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/observabilit original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: Optional[str] = None, - ): + ) -> Optional[HTTPException]: + """ + Transform error responses sent to clients. + + Return an HTTPException to replace the original error with a user-friendly message. + Return None to use the original exception. + + Example: + if isinstance(original_exception, litellm.ContextWindowExceededError): + return HTTPException( + status_code=400, + detail="Your prompt is too long. Please reduce the length and try again." + ) + return None # Use original exception + """ pass async def async_post_call_success_hook( @@ -339,3 +354,38 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ "usage": {} } ``` + +## Advanced - Transform Error Responses + +Transform technical API errors into user-friendly messages using `async_post_call_failure_hook`. Return an `HTTPException` to replace the original error, or `None` to use the original exception. + +```python +from litellm.integrations.custom_logger import CustomLogger +from fastapi import HTTPException +from typing import Optional +import litellm + +class MyErrorTransformer(CustomLogger): + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: Optional[str] = None, + ) -> Optional[HTTPException]: + if isinstance(original_exception, litellm.ContextWindowExceededError): + return HTTPException( + status_code=400, + detail="Your prompt is too long. Please reduce the length and try again." + ) + if isinstance(original_exception, litellm.RateLimitError): + return HTTPException( + status_code=429, + detail="Rate limit exceeded. Please try again in a moment." + ) + return None # Use original exception + +proxy_handler_instance = MyErrorTransformer() +``` + +**Result:** Clients receive `"Your prompt is too long..."` instead of `"ContextWindowExceededError: Prompt exceeds context window"`. diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 6771999cd35..4c4e6fa6342 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -32,6 +32,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.caching.caching import DualCache from opentelemetry.trace import Span as _Span @@ -348,7 +350,20 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: Optional[str] = None, - ): + ) -> Optional["HTTPException"]: + """ + Called after an LLM API call fails. Can return or raise HTTPException to transform error responses. + + Args: + - request_data: dict - The request data. + - original_exception: Exception - The original exception that occurred. + - user_api_key_dict: UserAPIKeyAuth - The user API key dictionary. + - traceback_str: Optional[str] - The traceback string. + + Returns: + - Optional[HTTPException]: Return an HTTPException to transform the error response sent to the client. + Return None to use the original exception. + """ pass async def async_post_call_success_hook( diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 2b9c4cdce6e..9c306acd2c6 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -2,7 +2,6 @@ Handles Authentication Errors """ -import asyncio from typing import TYPE_CHECKING, Any, Optional, Union from fastapi import HTTPException, Request, status @@ -90,15 +89,17 @@ class UserAPIKeyAuthExceptionHandler: api_key=api_key, request_route=route, ) - asyncio.create_task( - proxy_logging_obj.post_call_failure_hook( - request_data=request_data, - original_exception=e, - user_api_key_dict=user_api_key_dict, - error_type=ProxyErrorTypes.auth_error, - route=route, - ) + # Allow callbacks to transform the error response + transformed_exception = await proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=e, + user_api_key_dict=user_api_key_dict, + error_type=ProxyErrorTypes.auth_error, + route=route, ) + # Use transformed exception if callback returned one, otherwise use original + if transformed_exception is not None: + e = transformed_exception if isinstance(e, litellm.BudgetExceededError): raise ProxyException( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f798d218f1d..302dd5639ed 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -786,11 +786,15 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.exception( f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}" ) - await proxy_logging_obj.post_call_failure_hook( + # Allow callbacks to transform the error response + transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=self.data, ) + # Use transformed exception if callback returned one, otherwise use original + if transformed_exception is not None: + e = transformed_exception litellm_debug_info = getattr(e, "litellm_debug_info", "") verbose_proxy_logger.debug( "\033[1;31mAn error occurred: %s %s\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`", @@ -970,11 +974,15 @@ class ProxyBaseLLMRequestProcessing: str(e) ) ) - await proxy_logging_obj.post_call_failure_hook( + # Allow callbacks to transform the error response + transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_data, ) + # Use transformed exception if callback returned one, otherwise use original + if transformed_exception is not None: + e = transformed_exception verbose_proxy_logger.debug( f"\033[1;31mAn error occurred: {e}\n\n Debug this by setting `--debug`, e.g. `litellm --model gpt-3.5-turbo --debug`" ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec86139c73c..d595db4a2e0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1465,9 +1465,10 @@ class ProxyLogging: error_type: Optional[ProxyErrorTypes] = None, route: Optional[str] = None, traceback_str: Optional[str] = None, - ): + ) -> Optional[HTTPException]: """ Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body. + Callbacks can return or raise HTTPException to transform error responses sent to clients. Covers: 1. /chat/completions @@ -1481,6 +1482,10 @@ class ProxyLogging: - error_type: Optional[ProxyErrorTypes] - The error type. - route: Optional[str] - The route. - traceback_str: Optional[str] - The traceback string, sometimes upstream endpoints might need to send the upstream traceback. In which case we use this + + Returns: + - Optional[HTTPException]: If any callback returns or raises an HTTPException, the first one found is returned. + Otherwise, returns None and the original exception is used. """ ### ALERTING ### @@ -1522,6 +1527,9 @@ class ProxyLogging: original_exception=original_exception, ) + # Track the first HTTPException returned or raised by any callback + transformed_exception: Optional[HTTPException] = None + for callback in litellm.callbacks: try: _callback: Optional[CustomLogger] = None @@ -1532,19 +1540,31 @@ class ProxyLogging: else: _callback = callback # type: ignore if _callback is not None and isinstance(_callback, CustomLogger): - asyncio.create_task( - _callback.async_post_call_failure_hook( + try: + hook_result = await _callback.async_post_call_failure_hook( request_data=request_data, user_api_key_dict=user_api_key_dict, original_exception=original_exception, traceback_str=traceback_str, ) - ) + # If callback returned an HTTPException, use it (first one wins) + if isinstance(hook_result, HTTPException) and transformed_exception is None: + transformed_exception = hook_result + except HTTPException as e: + # If callback raised an HTTPException, use it (first one wins) + if transformed_exception is None: + transformed_exception = e + except Exception as e: + # Log non-HTTPException errors from callbacks but don't break the flow + verbose_proxy_logger.exception( + f"[Non-Blocking] Error in async_post_call_failure_hook callback: {e}" + ) except Exception as e: verbose_proxy_logger.exception( - f"[Non-Blocking] Error in post_call_failure_hook: {e}" + f"[Non-Blocking] Error setting up post_call_failure_hook callback: {e}" ) - return + + return transformed_exception def _is_proxy_only_llm_api_error( self, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 04aeddb8f28..fcc8c1f0f2e 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -278,8 +278,8 @@ async def test_proxy_admin_expired_key_from_cache(): mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() - # Mock post_call_failure_hook as async function - mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + # Mock post_call_failure_hook as async function returning None (no transformation) + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) # Mock prisma_client mock_prisma_client = MagicMock() diff --git a/tests/test_litellm/proxy/hooks/test_post_call_failure_hook_integration.py b/tests/test_litellm/proxy/hooks/test_post_call_failure_hook_integration.py new file mode 100644 index 00000000000..7223c2e1f02 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_post_call_failure_hook_integration.py @@ -0,0 +1,146 @@ +""" +Integration tests for async_post_call_failure_hook. + +Tests verify that the failure hook can transform error responses sent to clients, +similar to how async_post_call_success_hook can transform successful responses. +""" + +import os +import sys +import pytest +from typing import Optional +from unittest.mock import patch + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import HTTPException +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + + +class ErrorTransformerLogger(CustomLogger): + """Logger that transforms errors into user-friendly messages""" + + def __init__(self): + self.called = False + self.transformed_exception = None + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: Optional[str] = None, + ): + self.called = True + self.transformed_exception = HTTPException( + status_code=400, + detail="User-friendly error: Your request could not be processed." + ) + return self.transformed_exception + + +@pytest.mark.asyncio +async def test_failure_hook_transforms_error_response(): + """ + Test that async_post_call_failure_hook can transform error responses. + This mirrors how async_post_call_success_hook can transform successful responses. + """ + transformer = ErrorTransformerLogger() + + # Mock litellm.callbacks to include our transformer + with patch("litellm.callbacks", [transformer]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_exception = Exception("Technical error message") + request_data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Call the hook + result = await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) + + # Verify hook was called + assert transformer.called is True + + # Verify transformed exception is returned + assert result is not None + assert isinstance(result, HTTPException) + assert result.detail == "User-friendly error: Your request could not be processed." + + +@pytest.mark.asyncio +async def test_failure_hook_returns_none_when_no_transformation(): + """ + Test that hook returning None uses original exception. + """ + class NoOpLogger(CustomLogger): + def __init__(self): + self.called = False + + async def async_post_call_failure_hook(self, *args, **kwargs): + self.called = True + return None + + logger = NoOpLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_exception = Exception("Original error") + request_data = {"model": "test"} + user_api_key_dict = UserAPIKeyAuth(api_key="test") + + result = await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) + + # Should return None (original exception will be used) + assert result is None + assert logger.called is True + + +@pytest.mark.asyncio +async def test_failure_hook_handles_exceptions_gracefully(): + """ + Test that hook failures don't break the error flow. + """ + class FailingLogger(CustomLogger): + def __init__(self): + self.called = False + + async def async_post_call_failure_hook(self, *args, **kwargs): + self.called = True + raise RuntimeError("Hook crashed!") + + logger = FailingLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + original_exception = Exception("Original error") + request_data = {"model": "test"} + user_api_key_dict = UserAPIKeyAuth(api_key="test") + + # Should not raise, should handle gracefully + result = await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) + + # Should return None (original exception will be used) + assert result is None + assert logger.called is True + diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index b0e198d5e7e..0bb9924af82 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1148,7 +1148,7 @@ class TestBedrockLLMProxyRoute: mock_user_api_key_dict.allowed_model_region = None mock_proxy_logging_obj = Mock() - mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) endpoint = "model/test-model/converse" model = "test-model" @@ -1291,7 +1291,7 @@ class TestBedrockLLMProxyRoute: mock_user_api_key_dict = Mock() mock_user_api_key_dict.api_key = "test-key" mock_proxy_logging_obj = Mock() - mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() + mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) with patch( "litellm.passthrough.main.llm_passthrough_route", From 3a1baae45cb58a7050cd28a6f553010ed1c39768 Mon Sep 17 00:00:00 2001 From: prasadkona Date: Mon, 22 Dec 2025 11:59:20 -0800 Subject: [PATCH 024/330] feat(databricks): Add enhanced authentication, security features, and custom user-agent support - Add OAuth M2M (Machine-to-Machine) authentication via DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET - Add Databricks SDK auto-auth with automatic credential discovery - Add sensitive data redaction for secure logging (tokens, API keys, secrets) - Add custom user_agent parameter for partner attribution in Databricks telemetry - Support user_agent in LiteLLM Proxy via config.yaml litellm_params - Add 49 mocked unit tests for all new functionality - Add 13 E2E tests for real-world validation (skipped in CI) - Update documentation with new features and examples --- docs/my-website/docs/providers/databricks.md | 94 ++ .../llms/databricks/chat/transformation.py | 45 +- litellm/llms/databricks/common_utils.py | 311 ++++- litellm/llms/databricks/embed/handler.py | 12 + poetry.lock | 10 +- .../test_databricks_chat_transformation.py | 6 +- .../databricks/databricks_config.template.txt | 78 ++ .../llms/databricks/test_databricks_e2e.py | 1029 +++++++++++++++++ .../test_databricks_partner_integration.py | 662 +++++++++++ 9 files changed, 2218 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/llms/databricks/databricks_config.template.txt create mode 100644 tests/test_litellm/llms/databricks/test_databricks_e2e.py create mode 100644 tests/test_litellm/llms/databricks/test_databricks_partner_integration.py diff --git a/docs/my-website/docs/providers/databricks.md b/docs/my-website/docs/providers/databricks.md index 921b06a17b7..2791d55dff1 100644 --- a/docs/my-website/docs/providers/databricks.md +++ b/docs/my-website/docs/providers/databricks.md @@ -11,6 +11,99 @@ LiteLLM supports all models on Databricks ::: +## Authentication + +LiteLLM supports multiple authentication methods for Databricks, listed in order of preference: + +### OAuth M2M (Recommended for Production) + +OAuth Machine-to-Machine authentication using Service Principal credentials is the **recommended method for production** deployments per Databricks Partner requirements. + +```python +import os +from litellm import completion + +# Set OAuth credentials (Service Principal) +os.environ["DATABRICKS_CLIENT_ID"] = "your-service-principal-application-id" +os.environ["DATABRICKS_CLIENT_SECRET"] = "your-service-principal-secret" +os.environ["DATABRICKS_API_BASE"] = "https://adb-xxx.azuredatabricks.net/serving-endpoints" + +response = completion( + model="databricks/databricks-dbrx-instruct", + messages=[{"role": "user", "content": "Hello!"}], +) +``` + +### Personal Access Token (PAT) + +PAT authentication is supported for development and testing scenarios. + +```python +import os +from litellm import completion + +os.environ["DATABRICKS_API_KEY"] = "dapi..." # Your Personal Access Token +os.environ["DATABRICKS_API_BASE"] = "https://adb-xxx.azuredatabricks.net/serving-endpoints" + +response = completion( + model="databricks/databricks-dbrx-instruct", + messages=[{"role": "user", "content": "Hello!"}], +) +``` + +### Databricks SDK Authentication (Automatic) + +If no credentials are provided, LiteLLM will use the Databricks SDK for automatic authentication. This supports OAuth, Azure AD, and other unified auth methods configured in your environment. + +```python +from litellm import completion + +# No environment variables needed - uses Databricks SDK unified auth +# Requires: pip install databricks-sdk +response = completion( + model="databricks/databricks-dbrx-instruct", + messages=[{"role": "user", "content": "Hello!"}], +) +``` + +## Custom User-Agent for Partner Attribution + +If you're building a product on top of LiteLLM that integrates with Databricks, you can pass your own partner identifier for proper attribution in Databricks telemetry. + +The partner name will be prefixed to the LiteLLM user agent: + +```python +# Via parameter +response = completion( + model="databricks/databricks-dbrx-instruct", + messages=[{"role": "user", "content": "Hello!"}], + user_agent="mycompany/1.0.0", +) +# Resulting User-Agent: mycompany_litellm/1.79.1 + +# Via environment variable +os.environ["DATABRICKS_USER_AGENT"] = "mycompany/1.0.0" +# Resulting User-Agent: mycompany_litellm/1.79.1 +``` + +| Input | Resulting User-Agent | +|-------|---------------------| +| (none) | `litellm/1.79.1` | +| `mycompany/1.0.0` | `mycompany_litellm/1.79.1` | +| `partner_product/2.5.0` | `partner_product_litellm/1.79.1` | +| `acme` | `acme_litellm/1.79.1` | + +**Note:** The version from your custom user agent is ignored; LiteLLM's version is always used. + +## Security + +LiteLLM automatically redacts sensitive information (tokens, secrets, API keys) from all debug logs to prevent credential leakage. This includes: + +- Authorization headers +- API keys and tokens +- Client secrets +- Personal access tokens (PATs) + ## Usage @@ -51,6 +144,7 @@ response = completion( model: databricks/databricks-dbrx-instruct api_key: os.environ/DATABRICKS_API_KEY api_base: os.environ/DATABRICKS_API_BASE + user_agent: "mycompany/1.0.0" # Optional: for partner attribution ``` diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index ac3be0c3518..2b7f5dd5995 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -2,6 +2,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completions` """ +import os from typing import ( TYPE_CHECKING, Any, @@ -26,7 +27,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo _should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( - strip_name_from_message + strip_name_from_message, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.anthropic import AllAnthropicToolsValues @@ -124,12 +125,24 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: + # Check for custom user agent in optional_params or environment + # This allows partners building on LiteLLM to set their own telemetry + # Use pop() to remove these keys so they don't get sent to the API + custom_user_agent = ( + optional_params.pop("user_agent", None) + or optional_params.pop("databricks_user_agent", None) + or litellm_params.get("user_agent") + or os.getenv("LITELLM_USER_AGENT") + or os.getenv("DATABRICKS_USER_AGENT") + ) + api_base, headers = self.databricks_validate_environment( api_base=api_base, api_key=api_key, endpoint_type="chat_completions", custom_endpoint=False, headers=headers, + custom_user_agent=custom_user_agent, ) # Ensure Content-Type header is set headers["Content-Type"] = "application/json" @@ -173,9 +186,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): # Build DatabricksFunction explicitly to avoid parameter conflicts function_params: DatabricksFunction = { "name": tool["name"], - "parameters": cast(dict, tool.get("input_schema") or {}) + "parameters": cast(dict, tool.get("input_schema") or {}), } - + # Only add description if it exists description = tool.get("description") if description is not None: @@ -229,7 +242,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): Databricks supports Anthropic-style cache control for Claude models. Databricks ignores the cache_control flag with other models. """ - # TODO: Think about how to best design the request transformation so that + # TODO: Think about how to best design the request transformation so that # every request doesn't have to be transformed for to OpenAI and Anthropic request formats. return messages, tools @@ -347,15 +360,17 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): messages=new_messages, model=model, is_async=cast(Literal[False], False) ) - def _move_cache_control_into_string_content_block(self, message: AllMessageValues) -> AllMessageValues: + def _move_cache_control_into_string_content_block( + self, message: AllMessageValues + ) -> AllMessageValues: """ Moves message-level cache_control into a content block when content is a string. - + Transforms: {"role": "user", "content": "text", "cache_control": {...}} Into: {"role": "user", "content": [{"type": "text", "text": "text", "cache_control": {...}}]} - + This is required for Anthropic's prompt caching API when cache_control is specified at the message level but content is a simple string (not already an array of content blocks). """ @@ -371,7 +386,6 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): } ] return cast(AllMessageValues, transformed_message) - @staticmethod def extract_content_str( @@ -509,9 +523,9 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, tool_calls=choice["message"].get("tool_calls"), - provider_specific_fields={"citations": citations} - if citations is not None - else None, + provider_specific_fields=( + {"citations": citations} if citations is not None else None + ), ) if finish_reason is None: @@ -543,12 +557,15 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: - ## LOGGING + # Redact sensitive data before logging to prevent credential leakage + redacted_request_data = self.redact_sensitive_data(request_data) + + ## LOGGING - Never log actual API keys logging_obj.post_call( input=messages, - api_key=api_key, + api_key="[REDACTED]", original_response=raw_response.text, - additional_args={"complete_input_dict": request_data}, + additional_args={"complete_input_dict": redacted_request_data}, ) ## RESPONSE OBJECT diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 1353b5b13f6..608f29a03a7 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -1,4 +1,18 @@ -from typing import Literal, Optional, Tuple +""" +Databricks integration utilities for LiteLLM. + +This module provides authentication, telemetry, and security utilities +for the Databricks LLM provider integration. + +Authentication priority: +1. OAuth M2M (DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET) - Recommended for production +2. PAT (DATABRICKS_API_KEY) - Supported for development +3. Databricks SDK automatic auth - Fallback (uses unified auth) +""" + +import os +import re +from typing import Any, Dict, Literal, Optional, Tuple from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -8,17 +22,175 @@ class DatabricksException(BaseLLMException): class DatabricksBase: + """ + Base class for Databricks integration with authentication, + telemetry, and security utilities. + """ + + # Patterns to redact in logs + SENSITIVE_PATTERNS = [ + (re.compile(r"(Bearer\s+)[A-Za-z0-9\-_\.]+", re.IGNORECASE), r"\1[REDACTED]"), + (re.compile(r"(Authorization:\s*)[^\s,}]+", re.IGNORECASE), r"\1[REDACTED]"), + ( + re.compile(r'(api[_-]?key["\s:=]+)[^\s,}"\']+', re.IGNORECASE), + r"\1[REDACTED]", + ), + ( + re.compile(r'(client[_-]?secret["\s:=]+)[^\s,}"\']+', re.IGNORECASE), + r"\1[REDACTED]", + ), + (re.compile(r"(dapi[a-zA-Z0-9]{32,})", re.IGNORECASE), r"[REDACTED_PAT]"), + ( + re.compile(r'(access[_-]?token["\s:=]+)[^\s,}"\']+', re.IGNORECASE), + r"\1[REDACTED]", + ), + ] + + @classmethod + def redact_sensitive_data(cls, data: Any) -> Any: + """ + Redact sensitive information (tokens, secrets) from data before logging. + + Handles strings, dicts, and lists recursively. Keys containing sensitive + terms (authorization, api_key, token, secret, password, credential) are + fully redacted. + + Args: + data: String, dict, or other data structure to redact + + Returns: + Redacted version of the data safe for logging + """ + if data is None: + return None + + if isinstance(data, str): + result = data + for pattern, replacement in cls.SENSITIVE_PATTERNS: + result = pattern.sub(replacement, result) + return result + + if isinstance(data, dict): + redacted = {} + for key, value in data.items(): + lower_key = key.lower() + if any( + sensitive in lower_key + for sensitive in [ + "authorization", + "api_key", + "apikey", + "token", + "secret", + "password", + "credential", + ] + ): + redacted[key] = "[REDACTED]" + else: + redacted[key] = cls.redact_sensitive_data(value) + return redacted + + if isinstance(data, list): + return [cls.redact_sensitive_data(item) for item in data] + + return data + + @classmethod + def redact_headers_for_logging(cls, headers: Dict[str, str]) -> Dict[str, str]: + """ + Create a copy of headers with sensitive values redacted for safe logging. + + Shows first 8 characters of sensitive values for debugging purposes, + with the rest redacted. + + Args: + headers: HTTP headers dictionary + + Returns: + New dictionary with sensitive headers redacted + """ + if not headers: + return {} + + redacted = {} + sensitive_headers = { + "authorization", + "x-api-key", + "api-key", + "x-databricks-token", + } + + for key, value in headers.items(): + if key.lower() in sensitive_headers: + if len(value) > 10: + redacted[key] = f"{value[:8]}...[REDACTED]" + else: + redacted[key] = "[REDACTED]" + else: + redacted[key] = value + + return redacted + + @staticmethod + def _build_user_agent(custom_user_agent: Optional[str] = None) -> str: + """ + Build the User-Agent string for Databricks API calls. + + If a custom user agent is provided, the partner name (part before /) + is extracted and prefixed to the litellm user agent with an underscore. + The custom version is ignored; LiteLLM's version is always used. + + Args: + custom_user_agent: Optional custom user agent string (e.g., "mycompany/1.0.0") + + Returns: + User-Agent string in format: + - Default: "litellm/{version}" + - With custom: "{partner}_litellm/{version}" + + Examples: + - None -> "litellm/1.79.1" + - "mycompany/1.0.0" -> "mycompany_litellm/1.79.1" + - "partner_product/2.0.0" -> "partner_product_litellm/1.79.1" + - "acme" -> "acme_litellm/1.79.1" + """ + try: + from litellm._version import version + except Exception: + version = "0.0.0" + + if custom_user_agent: + custom_user_agent = custom_user_agent.strip() + + # Extract partner name (part before / if present) + if "/" in custom_user_agent: + partner_name = custom_user_agent.split("/")[0].strip() + else: + partner_name = custom_user_agent + + # Validate partner name: alphanumeric, underscore, hyphen only + if ( + partner_name + and partner_name.replace("_", "").replace("-", "").isalnum() + ): + return f"{partner_name}_litellm/{version}" + + # Default: just litellm + return f"litellm/{version}" + def _get_api_base(self, api_base: Optional[str]) -> str: + """ + Get the Databricks API base URL. + + If not provided, attempts to get it from the Databricks SDK. + """ if api_base is None: try: from databricks.sdk import WorkspaceClient databricks_client = WorkspaceClient() - - api_base = ( - api_base or f"{databricks_client.config.host}/serving-endpoints" - ) - + api_base = f"{databricks_client.config.host}/serving-endpoints" return api_base except ImportError: raise DatabricksException( @@ -30,12 +202,87 @@ class DatabricksBase: ) return api_base + def _get_oauth_m2m_token( + self, + api_base: str, + client_id: str, + client_secret: str, + ) -> str: + """ + Obtain an OAuth M2M access token using client credentials flow. + + This is the recommended authentication method for production integrations + per Databricks Partner requirements. + + Args: + api_base: Databricks workspace URL + client_id: OAuth client ID (Service Principal application ID) + client_secret: OAuth client secret + + Returns: + Access token string + + Raises: + DatabricksException: If token request fails + """ + import requests + + # Extract workspace URL from api_base + workspace_url = api_base.rstrip("/") + if "/serving-endpoints" in workspace_url: + workspace_url = workspace_url.replace("/serving-endpoints", "") + + token_url = f"{workspace_url}/oidc/v1/token" + + try: + response = requests.post( + token_url, + data={ + "grant_type": "client_credentials", + "scope": "all-apis", + }, + auth=(client_id, client_secret), + headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=30, + ) + except requests.RequestException as e: + raise DatabricksException( + status_code=500, + message=f"OAuth M2M token request failed: {str(e)}", + ) + + if response.status_code != 200: + raise DatabricksException( + status_code=response.status_code, + message=f"OAuth M2M token request failed: {response.text}", + ) + + token_data = response.json() + return token_data["access_token"] + def _get_databricks_credentials( self, api_key: Optional[str], api_base: Optional[str], headers: Optional[dict] ) -> Tuple[str, dict]: + """ + Get Databricks credentials using the Databricks SDK. + + Also registers LiteLLM as a partner for proper telemetry attribution + in Databricks system.access.audit table. + + Args: + api_key: Optional API key (PAT) + api_base: Optional API base URL + headers: Optional existing headers + + Returns: + Tuple of (api_base, headers) + """ headers = headers or {"Content-Type": "application/json"} try: - from databricks.sdk import WorkspaceClient + from databricks.sdk import WorkspaceClient, useragent + + # Register LiteLLM as partner for Databricks telemetry attribution + useragent.with_partner("litellm") databricks_client = WorkspaceClient() @@ -66,14 +313,53 @@ class DatabricksBase: endpoint_type: Literal["chat_completions", "embeddings"], custom_endpoint: Optional[bool], headers: Optional[dict], + custom_user_agent: Optional[str] = None, ) -> Tuple[str, dict]: - if api_key is None and not headers: # handle empty headers + """ + Validate and configure the Databricks environment. + + Authentication priority: + 1. OAuth M2M (DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET) - Recommended + 2. PAT (DATABRICKS_API_KEY) - Supported for development + 3. Databricks SDK automatic auth - Fallback (uses unified auth) + + Args: + api_key: Personal access token (PAT) + api_base: Databricks workspace URL with /serving-endpoints + endpoint_type: Type of endpoint (chat_completions or embeddings) + custom_endpoint: Whether using a custom endpoint URL + headers: Existing headers dict + custom_user_agent: Optional custom user agent to prefix + + Returns: + Tuple of (api_base, headers) with authentication configured + """ + from litellm._logging import verbose_logger + + # Check for OAuth M2M credentials (recommended for production) + client_id = os.getenv("DATABRICKS_CLIENT_ID") + client_secret = os.getenv("DATABRICKS_CLIENT_SECRET") + + # Determine api_base first + if api_base is None: + api_base = os.getenv("DATABRICKS_API_BASE") + + if client_id and client_secret and api_base: + # Use OAuth M2M flow (preferred for production) + verbose_logger.debug("Using OAuth M2M authentication for Databricks") + access_token = self._get_oauth_m2m_token(api_base, client_id, client_secret) + headers = headers or {} + headers["Authorization"] = f"Bearer {access_token}" + headers["Content-Type"] = "application/json" + elif api_key is None and not headers: if custom_endpoint is True: raise DatabricksException( status_code=400, message="Missing API Key - A call is being made to LLM Provider but no key is set either in the environment variables ({LLM_PROVIDER}_API_KEY) or via params", ) else: + # Fallback to Databricks SDK (registers partner telemetry) + verbose_logger.debug("Using Databricks SDK for authentication") api_base, headers = self._get_databricks_credentials( api_base=api_base, api_key=api_key, headers=headers ) @@ -101,8 +387,17 @@ class DatabricksBase: if api_key is not None: headers["Authorization"] = f"Bearer {api_key}" + # Set User-Agent with optional custom prefix + headers["User-Agent"] = self._build_user_agent(custom_user_agent) + + # Debug logging with redaction (never log actual tokens) + verbose_logger.debug( + f"Databricks request headers: {self.redact_headers_for_logging(headers)}" + ) + if endpoint_type == "chat_completions" and custom_endpoint is not True: api_base = "{}/chat/completions".format(api_base) elif endpoint_type == "embeddings" and custom_endpoint is not True: api_base = "{}/embeddings".format(api_base) + return api_base, headers diff --git a/litellm/llms/databricks/embed/handler.py b/litellm/llms/databricks/embed/handler.py index 2eabcdbc866..227824f72d0 100644 --- a/litellm/llms/databricks/embed/handler.py +++ b/litellm/llms/databricks/embed/handler.py @@ -2,6 +2,7 @@ Calling logic for Databricks embeddings """ +import os from typing import Optional from litellm.utils import EmbeddingResponse @@ -26,12 +27,23 @@ class DatabricksEmbeddingHandler(OpenAILikeEmbeddingHandler, DatabricksBase): custom_endpoint: Optional[bool] = None, headers: Optional[dict] = None, ) -> EmbeddingResponse: + # Check for custom user agent in optional_params or environment + # This allows partners building on LiteLLM to set their own telemetry + # Use pop() to remove these keys so they don't get sent to the API + custom_user_agent = ( + optional_params.pop("user_agent", None) + or optional_params.pop("databricks_user_agent", None) + or os.getenv("LITELLM_USER_AGENT") + or os.getenv("DATABRICKS_USER_AGENT") + ) + api_base, headers = self.databricks_validate_environment( api_base=api_base, api_key=api_key, endpoint_type="embeddings", custom_endpoint=custom_endpoint, headers=headers, + custom_user_agent=custom_user_agent, ) return super().embedding( model=model, diff --git a/poetry.lock b/poetry.lock index 5313167def4..4ae5cf01079 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "aiofiles" @@ -3138,15 +3138,15 @@ openai = ["openai (>=0.27.8)"] [[package]] name = "litellm-enterprise" -version = "0.1.25" +version = "0.1.27" description = "Package for LiteLLM Enterprise features" optional = true python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8" groups = ["main"] markers = "extra == \"proxy\"" files = [ - {file = "litellm_enterprise-0.1.25-py3-none-any.whl", hash = "sha256:80c8f1996846453ad309e74cd6d2659d9508320370df5d462d34326b06401c4d"}, - {file = "litellm_enterprise-0.1.25.tar.gz", hash = "sha256:1c82178b8e2c85f47b31910fd103a322b46d6caea44cd7a8c80b00fdcfeacd22"}, + {file = "litellm_enterprise-0.1.27-py3-none-any.whl", hash = "sha256:41b9d41d04123f492060a742091006dc1d182b54ce3a1c0e18ee75d623c63e91"}, + {file = "litellm_enterprise-0.1.27.tar.gz", hash = "sha256:aa40c87f7c8df64beb79e75f71e1b5c0a458350efa68527e3491e6f27f2cbd57"}, ] [[package]] @@ -8051,4 +8051,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "a102d24777f1c438dcf15055abeef385722a9603cb3c3d3643c86190b7534c47" +content-hash = "996152bfbb1d7870a4b6f1837d7ff556320d6e817198ecec4198ed99539e848b" diff --git a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py index a14683fac17..f437b8405f7 100644 --- a/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py @@ -94,12 +94,13 @@ def test_transform_choices_without_signature(): assert thinking_block["type"] == "thinking" assert thinking_block["thinking"] == "i'm thinking without signature." + def test_convert_anthropic_tool_to_databricks_tool_with_description(): config = DatabricksConfig() anthropic_tool = { "name": "test_tool", "description": "test description", - "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}} + "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}}, } databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool) @@ -113,7 +114,7 @@ def test_convert_anthropic_tool_to_databricks_tool_without_description(): config = DatabricksConfig() anthropic_tool = { "name": "test_tool", - "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}} + "input_schema": {"type": "object", "properties": {"test": {"type": "string"}}}, } databricks_tool = config.convert_anthropic_tool_to_databricks_tool(anthropic_tool) @@ -122,6 +123,7 @@ def test_convert_anthropic_tool_to_databricks_tool_without_description(): assert databricks_tool["type"] == "function" assert databricks_tool["function"].get("description") is None + def test_transform_choices_with_citations(): config = DatabricksConfig() databricks_choices = [ diff --git a/tests/test_litellm/llms/databricks/databricks_config.template.txt b/tests/test_litellm/llms/databricks/databricks_config.template.txt new file mode 100644 index 00000000000..7352fdbc773 --- /dev/null +++ b/tests/test_litellm/llms/databricks/databricks_config.template.txt @@ -0,0 +1,78 @@ +# Databricks Configuration Template for LiteLLM Testing +# ===================================================== +# +# Copy this file to your preferred location and fill in your credentials: +# cp databricks_config.template.txt /path/to/databricks_config.txt +# +# Then update the CONFIG_FILE path in test_databricks_integration.py +# +# Lines starting with # are comments and will be ignored +# Only lines with KEY=VALUE format (where VALUE is not empty) will be read + +# ============================================================================== +# DATABRICKS WORKSPACE CONFIGURATION (Required) +# ============================================================================== + +# Your Databricks workspace URL (without /serving-endpoints suffix) +# Example: https://adb-1234567890123456.7.azuredatabricks.net +DATABRICKS_HOST= + +# API Base URL for serving endpoints (usually {host}/serving-endpoints) +# Example: https://adb-1234567890123456.7.azuredatabricks.net/serving-endpoints +DATABRICKS_API_BASE= + +# ============================================================================== +# AUTHENTICATION METHOD 1: OAuth M2M (Recommended for Production) +# Use Service Principal credentials +# ============================================================================== + +# Service Principal Application/Client ID +# Example: 12345678-1234-1234-1234-123456789012 +DATABRICKS_CLIENT_ID= + +# Service Principal Secret +# Example: your-client-secret-value +DATABRICKS_CLIENT_SECRET= + +# ============================================================================== +# AUTHENTICATION METHOD 2: Personal Access Token (PAT) +# For development and testing +# ============================================================================== + +# Personal Access Token (starts with 'dapi') +# Example: dapi_your_token_here +DATABRICKS_API_KEY= + +# ============================================================================== +# MODEL CONFIGURATION +# ============================================================================== + +# Model to use for testing chat completions +# Example: databricks-gpt-oss-120b, databricks-meta-llama-3-1-70b-instruct +TEST_CHAT_MODEL=databricks-gpt-oss-120b + +# Model to use for testing embeddings (optional) +# Example: databricks-bge-large-en +TEST_EMBEDDING_MODEL=databricks-bge-large-en + +# ============================================================================== +# OPTIONAL: Custom User-Agent for Partner Attribution Testing +# ============================================================================== + +# Custom user agent string to test partner attribution +# Example: mycompany/1.0.0 +# This will result in User-Agent: mycompany_litellm/{version} +# Leave empty to use default: litellm/{version} +CUSTOM_USER_AGENT= + +# ============================================================================== +# TEST SETTINGS +# ============================================================================== + +# Which authentication method to test: oauth, pat, sdk, or all +# oauth = Use DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET +# pat = Use DATABRICKS_API_KEY +# sdk = Use Databricks SDK automatic authentication (~/.databrickscfg) +# all = Test all three methods (oauth, pat, sdk) in sequence +TEST_AUTH_METHOD=pat + diff --git a/tests/test_litellm/llms/databricks/test_databricks_e2e.py b/tests/test_litellm/llms/databricks/test_databricks_e2e.py new file mode 100644 index 00000000000..669f9e94639 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_e2e.py @@ -0,0 +1,1029 @@ +""" +End-to-End Tests for Databricks LiteLLM Integration +==================================================== + +⚠️ WARNING: These tests require REAL Databricks credentials and make ACTUAL API calls. + They are NOT suitable for automated CI/CD pipelines. + +For unit tests that use mocks and don't require credentials, see: + test_databricks_partner_integration.py + +Purpose: + - Validate actual API connectivity with Databricks + - Test all authentication methods (OAuth M2M, PAT, SDK) + - Verify User-Agent strings appear correctly in Databricks audit logs + - Test chat completions and embeddings with real models + - Test different SDK integration methods with custom user agents + +LiteLLM Integration Tests: + This test file includes tests for different ways of calling Databricks via LiteLLM: + + 1. LiteLLM SDK Direct - Using litellm.completion() with user_agent parameter + 2. LangChain + LiteLLM - Using ChatLiteLLM wrapper (requires langchain-community) + 3. LiteLLM Async - Using litellm.acompletion() async API + 4. LiteLLM Streaming - Using litellm.completion() with stream=True + 5. LiteLLM Embedding - Using litellm.embedding() with user_agent parameter + + All tests use the CUSTOM_USER_AGENT value from the config file and call + Databricks endpoints through LiteLLM's unified interface. + +Prerequisites: + - Valid Databricks workspace access + - Configured credentials (OAuth Service Principal, PAT, or Databricks CLI) + - Access to serving endpoints (e.g., databricks-gpt-oss-120b) + +Optional Dependencies (for LiteLLM integration tests): + - pip install langchain-litellm # For LangChain tests (recommended) + +Setup: + 1. Copy the template to create your config file: + cp databricks_config.template.txt ~/.databricks_litellm_config.txt + + 2. Edit the config file with your Databricks credentials: + - DATABRICKS_API_BASE (required) + - DATABRICKS_HOST (required for Databricks SDK tests) + - DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET (for OAuth) + - DATABRICKS_API_KEY (for PAT) + - CUSTOM_USER_AGENT (for partner attribution tests) + + 3. Optionally set a custom config path: + export DATABRICKS_TEST_CONFIG=/path/to/your/config.txt + +Run with: + cd /path/to/litellm + python tests/test_litellm/llms/databricks/test_databricks_e2e.py + +Config Options: + TEST_AUTH_METHOD=oauth # Test OAuth M2M authentication + TEST_AUTH_METHOD=pat # Test Personal Access Token + TEST_AUTH_METHOD=sdk # Test Databricks SDK (~/.databrickscfg) + TEST_AUTH_METHOD=all # Test all three methods sequentially +""" + +import os +import sys + +import pytest + +# Skip all tests in this module during unit test runs (make test-unit) +# These are E2E tests that require real Databricks credentials +pytestmark = pytest.mark.skip( + reason="E2E tests require real Databricks credentials. Run directly with: " + "python tests/test_litellm/llms/databricks/test_databricks_e2e.py" +) + +# Add the litellm package to path +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +) + +# Config file path - can be overridden with DATABRICKS_TEST_CONFIG env var +DEFAULT_CONFIG_PATH = os.path.expanduser("~/.databricks_litellm_config.txt") +CONFIG_FILE = os.environ.get("DATABRICKS_TEST_CONFIG", DEFAULT_CONFIG_PATH) + + +def load_config(config_file: str) -> dict: + """Load configuration from file.""" + config = {} + + template_path = os.path.join( + os.path.dirname(__file__), "databricks_config.template.txt" + ) + + if not os.path.exists(config_file): + raise FileNotFoundError( + f"Config file not found: {config_file}\n\n" + f"To set up:\n" + f" 1. Copy the template:\n" + f" cp {template_path} {config_file}\n\n" + f" 2. Edit {config_file} with your Databricks credentials\n\n" + f" 3. Or set a custom path:\n" + f" export DATABRICKS_TEST_CONFIG=/your/path/config.txt" + ) + + with open(config_file, "r") as f: + for line in f: + line = line.strip() + # Skip comments and empty lines + if not line or line.startswith("#"): + continue + + # Parse KEY=VALUE + if "=" in line: + key, value = line.split("=", 1) + key = key.strip() + value = value.strip() + if value: # Only set if value is not empty + config[key] = value + + return config + + +def setup_environment(config: dict, auth_method: str): + """Set up environment variables based on auth method.""" + # Clear any existing Databricks env vars (including SDK-specific ones) + for var in [ + "DATABRICKS_API_KEY", + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CLIENT_SECRET", + "DATABRICKS_API_BASE", + "DATABRICKS_USER_AGENT", + "LITELLM_USER_AGENT", + "DATABRICKS_TOKEN", + "DATABRICKS_HOST", + ]: # Added SDK env vars + os.environ.pop(var, None) + + # Set auth based on method + if auth_method == "oauth": + if ( + "DATABRICKS_CLIENT_ID" not in config + or "DATABRICKS_CLIENT_SECRET" not in config + ): + raise ValueError( + "OAuth auth requires DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET" + ) + # For OAuth, set the API base + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + os.environ["DATABRICKS_CLIENT_ID"] = config["DATABRICKS_CLIENT_ID"] + os.environ["DATABRICKS_CLIENT_SECRET"] = config["DATABRICKS_CLIENT_SECRET"] + print(" Auth method: OAuth M2M (Service Principal)") + + elif auth_method == "pat": + if "DATABRICKS_API_KEY" not in config: + raise ValueError("PAT auth requires DATABRICKS_API_KEY") + # For PAT, set the API base + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + os.environ["DATABRICKS_API_KEY"] = config["DATABRICKS_API_KEY"] + print(" Auth method: Personal Access Token (PAT)") + + elif auth_method == "sdk": + # For SDK mode, don't set any env vars - let SDK use ~/.databrickscfg + # But we still need to pass api_base to litellm, so set it if provided + if "DATABRICKS_API_BASE" in config: + os.environ["DATABRICKS_API_BASE"] = config["DATABRICKS_API_BASE"] + print(" Auth method: Databricks SDK (automatic from ~/.databrickscfg)") + + else: + raise ValueError(f"Unknown auth method: {auth_method}") + + # Set custom user agent if provided + if "CUSTOM_USER_AGENT" in config: + os.environ["DATABRICKS_USER_AGENT"] = config["CUSTOM_USER_AGENT"] + print(f" Custom User-Agent: {config['CUSTOM_USER_AGENT']}") + + +def test_user_agent_building(): + """Test User-Agent string building.""" + print("\n" + "=" * 60) + print("TEST: User-Agent Building") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + # Test 1: Default + ua = DatabricksBase._build_user_agent(None) + print(f" Default: {ua}") + assert ua.startswith("litellm/"), f"Expected litellm/, got {ua}" + print(" ✓ Default user agent works") + + # Test 2: With partner + ua = DatabricksBase._build_user_agent("mycompany/1.0.0") + print(f" With partner: {ua}") + assert ua.startswith("mycompany_litellm/"), f"Expected mycompany_litellm/, got {ua}" + print(" ✓ Partner prefixing works") + + # Test 3: Partner without version + ua = DatabricksBase._build_user_agent("acme") + print(f" Without version: {ua}") + assert ua.startswith("acme_litellm/"), f"Expected acme_litellm/, got {ua}" + print(" ✓ Partner without version works") + + print(" ✓ All user agent tests passed!") + + +def test_token_redaction(): + """Test sensitive data redaction.""" + print("\n" + "=" * 60) + print("TEST: Token Redaction") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + # Test header redaction + headers = { + "Authorization": "Bearer dapi123456789abcdef", + "Content-Type": "application/json", + } + redacted = DatabricksBase.redact_headers_for_logging(headers) + print(f" Original: Authorization: Bearer dapi123456789abcdef") + print(f" Redacted: Authorization: {redacted['Authorization']}") + assert "[REDACTED]" in redacted["Authorization"] + assert redacted["Content-Type"] == "application/json" + print(" ✓ Header redaction works") + + # Test dict redaction + data = {"api_key": "secret123", "model": "dbrx"} + redacted = DatabricksBase.redact_sensitive_data(data) + assert redacted["api_key"] == "[REDACTED]" + assert redacted["model"] == "dbrx" + print(" ✓ Dict redaction works") + + # Test PAT redaction + text = "Token: dapi_fake_test_token_for_testing" + redacted = DatabricksBase.redact_sensitive_data(text) + assert "dapi_fake_test" not in redacted + print(" ✓ PAT string redaction works") + + print(" ✓ All redaction tests passed!") + + +def test_chat_completion(config: dict): + """Test chat completion with Databricks.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion") + print("=" * 60) + + import litellm + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" API Base: {os.environ.get('DATABRICKS_API_BASE', 'Not set')}") + + try: + response = litellm.completion( + model=full_model, + messages=[ + { + "role": "user", + "content": "Say 'Hello, LiteLLM test!' in exactly those words.", + } + ], + max_tokens=50, + temperature=0.1, + ) + + content = response.choices[0].message.content + print(f" Response: {content[:100]}...") + print(f" Model returned: {response.model}") + print(f" Usage: {response.usage}") + print(" ✓ Chat completion test passed!") + return True + + except Exception as e: + print(f" ✗ Chat completion failed: {e}") + return False + + +def test_chat_completion_default_user_agent(config: dict): + """Test chat completion with default user agent (no custom agent).""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with DEFAULT User-Agent") + print("=" * 60) + + import litellm + + # Clear any custom user agent from environment + saved_user_agent = os.environ.pop("DATABRICKS_USER_AGENT", None) + saved_litellm_ua = os.environ.pop("LITELLM_USER_AGENT", None) + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" Expected User-Agent: litellm/{version}") + print(f" (No custom user agent set)") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'default' only."}], + max_tokens=10, + # Note: NOT passing user_agent parameter + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Default user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is 'litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Default user-agent test failed: {e}") + return False + + finally: + # Restore environment variables + if saved_user_agent: + os.environ["DATABRICKS_USER_AGENT"] = saved_user_agent + if saved_litellm_ua: + os.environ["LITELLM_USER_AGENT"] = saved_litellm_ua + + +def test_chat_completion_with_custom_user_agent(config: dict): + """Test chat completion with custom user agent passed as parameter.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with Custom User-Agent (parameter)") + print("=" * 60) + + import litellm + + # Clear any env user agent to ensure parameter takes precedence + saved_user_agent = os.environ.pop("DATABRICKS_USER_AGENT", None) + saved_litellm_ua = os.environ.pop("LITELLM_USER_AGENT", None) + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" Custom User-Agent param: testpartner/2.0.0") + print(f" Expected User-Agent: testpartner_litellm/{version}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'test' only."}], + max_tokens=10, + user_agent="testpartner/2.0.0", # This should result in testpartner_litellm/{version} + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Custom user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is 'testpartner_litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Custom user-agent test failed: {e}") + return False + + finally: + # Restore environment variables + if saved_user_agent: + os.environ["DATABRICKS_USER_AGENT"] = saved_user_agent + if saved_litellm_ua: + os.environ["LITELLM_USER_AGENT"] = saved_litellm_ua + + +def test_chat_completion_with_env_user_agent(config: dict): + """Test chat completion with user agent set via environment variable.""" + print("\n" + "=" * 60) + print("TEST: Chat Completion with User-Agent from ENV VAR") + print("=" * 60) + + import litellm + + # Set a specific user agent via environment + test_partner = "envpartner" + os.environ["DATABRICKS_USER_AGENT"] = test_partner + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + print(f" DATABRICKS_USER_AGENT env var: {test_partner}") + print(f" Expected User-Agent: {test_partner}_litellm/{version}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'env' only."}], + max_tokens=10, + # Note: NOT passing user_agent parameter - should use env var + ) + + print(f" Response: {response.choices[0].message.content}") + print(" ✓ Env var user-agent test passed!") + print( + f" Note: Check Databricks Query History to verify User-Agent is '{test_partner}_litellm/{version}'" + ) + return True + + except Exception as e: + print(f" ✗ Env var user-agent test failed: {e}") + return False + + finally: + # Clean up + os.environ.pop("DATABRICKS_USER_AGENT", None) + + +def test_embedding(config: dict): + """Test embeddings with Databricks.""" + print("\n" + "=" * 60) + print("TEST: Embeddings") + print("=" * 60) + + import litellm + + model = config.get("TEST_EMBEDDING_MODEL", "databricks-bge-large-en") + full_model = f"databricks/{model}" + + print(f" Model: {full_model}") + + try: + response = litellm.embedding( + model=full_model, + input=["Hello, world!"], + ) + + # Handle both object and dict response formats + if hasattr(response, "data"): + data = response.data + else: + data = response.get("data", []) + + if data: + first_item = data[0] + if hasattr(first_item, "embedding"): + embedding = first_item.embedding + else: + embedding = first_item.get("embedding", []) + + print(f" Embedding dimensions: {len(embedding)}") + print(f" First 5 values: {embedding[:5]}") + print(" ✓ Embedding test passed!") + return True + else: + print(" ✗ Embedding test failed: No data in response") + return False + + except Exception as e: + print(f" ✗ Embedding test failed: {e}") + print(" (This is expected if embedding model is not available)") + return False + + +def test_oauth_token_retrieval(config: dict): + """Test OAuth M2M token retrieval.""" + print("\n" + "=" * 60) + print("TEST: OAuth M2M Token Retrieval") + print("=" * 60) + + if "DATABRICKS_CLIENT_ID" not in config or "DATABRICKS_CLIENT_SECRET" not in config: + print(" Skipped: OAuth credentials not configured") + return None + + from litellm.llms.databricks.common_utils import DatabricksBase + + try: + db = DatabricksBase() + token = db._get_oauth_m2m_token( + api_base=config["DATABRICKS_API_BASE"], + client_id=config["DATABRICKS_CLIENT_ID"], + client_secret=config["DATABRICKS_CLIENT_SECRET"], + ) + + # Redact token for display + redacted_token = ( + f"{token[:10]}...[REDACTED]" if len(token) > 10 else "[REDACTED]" + ) + print(f" Token obtained: {redacted_token}") + print(" ✓ OAuth M2M token retrieval passed!") + return True + + except Exception as e: + print(f" ✗ OAuth token retrieval failed: {e}") + return False + + +# ============================================================================== +# SDK INTEGRATION TESTS - Different ways of calling Databricks via LiteLLM +# ============================================================================== + + +def test_litellm_sdk_with_config_user_agent(config: dict): + """ + Test 1: LiteLLM SDK with custom user agent from config file. + + This test uses the LiteLLM SDK directly with the CUSTOM_USER_AGENT + specified in the databricks config file. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM SDK with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + try: + from litellm._version import version + except Exception: + version = "unknown" + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + response = litellm.completion( + model=full_model, + messages=[{"role": "user", "content": "Say 'LiteLLM SDK test' only."}], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, # Use config user agent + ) + + content = response.choices[0].message.content + print(f" Response: {content}") + print(" ✓ LiteLLM SDK with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM SDK test failed: {e}") + return False + + +def test_langchain_litellm_with_user_agent(config: dict): + """ + Test 2: LangChain with LiteLLM integration. + + This test uses LangChain's ChatLiteLLM wrapper to call Databricks + with custom user agent from config. + + Requires: pip install langchain-litellm (recommended) + or: pip install langchain langchain-community (deprecated) + """ + print("\n" + "=" * 60) + print("TEST: LangChain + LiteLLM with Config User-Agent") + print("=" * 60) + + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + # Try the new langchain-litellm package first, fall back to deprecated import + ChatLiteLLM = None + HumanMessage = None + + try: + from langchain_litellm import ChatLiteLLM + from langchain_core.messages import HumanMessage + + print(" Using: langchain-litellm package (recommended)") + except ImportError: + try: + # Fall back to deprecated import + import warnings + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=DeprecationWarning) + from langchain_community.chat_models import ChatLiteLLM + from langchain_core.messages import HumanMessage + print( + " Using: langchain-community (deprecated, consider: pip install langchain-litellm)" + ) + except ImportError: + print(" Skipped: langchain-litellm not installed") + print(" Install with: pip install langchain-litellm") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + # Set user agent via environment for LangChain integration + os.environ["DATABRICKS_USER_AGENT"] = custom_ua + + chat = ChatLiteLLM( + model=full_model, + max_tokens=20, + temperature=0.1, + ) + + messages = [HumanMessage(content="Say 'LangChain test' only.")] + response = chat.invoke(messages) + + content = response.content + print(f" Response: {content}") + print(" ✓ LangChain + LiteLLM with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LangChain + LiteLLM test failed: {e}") + import traceback + + traceback.print_exc() + return False + + finally: + # Clean up env var + os.environ.pop("DATABRICKS_USER_AGENT", None) + + +def test_litellm_async_completion(config: dict): + """ + Test 3: LiteLLM Async Completion API with custom User-Agent. + + This test uses LiteLLM's async completion API (acompletion) to call + Databricks with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Async Completion with Config User-Agent") + print("=" * 60) + + import asyncio + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + async def run_async_completion(): + response = await litellm.acompletion( + model=full_model, + messages=[{"role": "user", "content": "Say 'LiteLLM async test' only."}], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, + ) + return response + + try: + response = asyncio.run(run_async_completion()) + + content = response.choices[0].message.content + print(f" Response: {content}") + print(" ✓ LiteLLM async completion with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM async completion test failed: {e}") + import traceback + + traceback.print_exc() + return False + + +def test_litellm_streaming_completion(config: dict): + """ + Test 4: LiteLLM Streaming Completion with custom User-Agent. + + This test uses LiteLLM's streaming completion API to call + Databricks with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Streaming Completion with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_CHAT_MODEL", "databricks-gpt-oss-120b") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + # Use streaming completion + response = litellm.completion( + model=full_model, + messages=[ + {"role": "user", "content": "Say 'LiteLLM streaming test' only."} + ], + max_tokens=20, + temperature=0.1, + user_agent=custom_ua, + stream=True, + ) + + # Collect streamed content + collected_content = "" + for chunk in response: + if chunk.choices and chunk.choices[0].delta.content: + collected_content += chunk.choices[0].delta.content + + print(f" Response (streamed): {collected_content}") + print(" ✓ LiteLLM streaming completion with config user-agent test passed!") + return True + + except Exception as e: + print(f" ✗ LiteLLM streaming completion test failed: {e}") + import traceback + + traceback.print_exc() + return False + + +def test_litellm_embedding_with_user_agent(config: dict): + """ + Test 5: LiteLLM Embedding API with custom User-Agent. + + This test uses LiteLLM's embedding API to call Databricks + with custom user agent from config. + """ + print("\n" + "=" * 60) + print("TEST: LiteLLM Embedding with Config User-Agent") + print("=" * 60) + + import litellm + from litellm.llms.databricks.common_utils import DatabricksBase + + custom_ua = config.get("CUSTOM_USER_AGENT") + if not custom_ua: + print(" Skipped: CUSTOM_USER_AGENT not set in config") + return None + + model = config.get("TEST_EMBEDDING_MODEL", "databricks-bge-large-en") + full_model = f"databricks/{model}" + + # Build and display the final User-Agent that will be sent + final_user_agent = DatabricksBase._build_user_agent(custom_ua) + + print(f" Model: {full_model}") + print(f" Custom User-Agent from config: {custom_ua}") + print(f" >>> Final User-Agent sent: {final_user_agent}") + + try: + response = litellm.embedding( + model=full_model, + input=["Hello, this is a LiteLLM embedding test with custom user agent!"], + user_agent=custom_ua, + ) + + # Handle both object and dict response formats + if hasattr(response, "data"): + data = response.data + else: + data = response.get("data", []) + + if data: + first_item = data[0] + if hasattr(first_item, "embedding"): + embedding = first_item.embedding + else: + embedding = first_item.get("embedding", []) + + print(f" Embedding dimensions: {len(embedding)}") + print(f" First 3 values: {embedding[:3]}") + print(" ✓ LiteLLM embedding with config user-agent test passed!") + return True + else: + print(" ✗ LiteLLM embedding test failed: No data in response") + return False + + except Exception as e: + print(f" ✗ LiteLLM embedding test failed: {e}") + print(" (This may fail if embedding model is not available)") + import traceback + + traceback.print_exc() + return False + + +def run_integration_tests_for_auth_method(config: dict, auth_method: str) -> list: + """Run integration tests for a specific auth method. Returns list of (name, result) tuples.""" + results = [] + + print("\n" + "=" * 60) + print(f"INTEGRATION TESTS - {auth_method.upper()} Authentication") + print("=" * 60) + + # Setup environment for this auth method + try: + setup_environment(config, auth_method) + except ValueError as e: + print(f" ✗ Setup failed: {e}") + return [(f"[{auth_method.upper()}] Setup", False)] + + # Test OAuth token retrieval (only for oauth method) + if auth_method == "oauth": + results.append( + ( + f"[{auth_method.upper()}] OAuth Token Retrieval", + test_oauth_token_retrieval(config), + ) + ) + + # Test chat completion + results.append( + (f"[{auth_method.upper()}] Chat Completion", test_chat_completion(config)) + ) + + # Test embeddings + results.append((f"[{auth_method.upper()}] Embeddings", test_embedding(config))) + + return results + + +def main(): + print("=" * 60) + print("DATABRICKS LITELLM INTEGRATION TESTS") + print("=" * 60) + + # Load config + print(f"\nLoading config from: {CONFIG_FILE}") + try: + config = load_config(CONFIG_FILE) + print(f" Loaded {len(config)} configuration values") + except FileNotFoundError as e: + print(f"\nERROR: {e}") + return 1 + + # Validate required config + if "DATABRICKS_API_BASE" not in config: + print("\nERROR: DATABRICKS_API_BASE is required in config file") + return 1 + + auth_method = config.get("TEST_AUTH_METHOD", "pat").lower() + print(f"\nTest Configuration:") + print(f" API Base: {config['DATABRICKS_API_BASE']}") + print(f" Auth Method: {auth_method}") + + # Run unit tests (no credentials needed) + print("\n" + "=" * 60) + print("UNIT TESTS (No credentials needed)") + print("=" * 60) + + test_user_agent_building() + test_token_redaction() + + all_results = [] + + # Determine which auth methods to test + if auth_method == "all": + auth_methods_to_test = ["oauth", "pat", "sdk"] + print("\n" + "#" * 60) + print("# TESTING ALL AUTHENTICATION METHODS") + print("#" * 60) + else: + auth_methods_to_test = [auth_method] + + # Run integration tests for each auth method + for method in auth_methods_to_test: + results = run_integration_tests_for_auth_method(config, method) + all_results.extend(results) + + # Run User-Agent tests (only once, using the last auth method or 'pat' for 'all') + print("\n" + "-" * 60) + print("USER-AGENT INTEGRATION TESTS") + print("-" * 60) + + # Setup environment for user-agent tests (use 'pat' as it's simplest) + if auth_method == "all": + setup_environment(config, "pat") + + # Test 1: Default user agent (no custom agent set) + all_results.append( + ( + "Chat with DEFAULT User-Agent", + test_chat_completion_default_user_agent(config), + ) + ) + + # Test 2: Custom user agent passed as parameter + all_results.append( + ( + "Chat with Custom User-Agent (param)", + test_chat_completion_with_custom_user_agent(config), + ) + ) + + # Test 3: User agent from environment variable + all_results.append( + ( + "Chat with User-Agent from ENV", + test_chat_completion_with_env_user_agent(config), + ) + ) + + # Run SDK Integration Tests with different calling methods + print("\n" + "#" * 60) + print("# SDK INTEGRATION TESTS - DIFFERENT CALLING METHODS") + print("# Using CUSTOM_USER_AGENT from config file") + print("#" * 60) + + # Setup environment for SDK tests (use 'pat' as it's most compatible) + setup_environment(config, "pat") + + # Test 1: LiteLLM SDK with config user agent + all_results.append( + ( + "LiteLLM SDK with Config User-Agent", + test_litellm_sdk_with_config_user_agent(config), + ) + ) + + # Test 2: LangChain + LiteLLM with config user agent + all_results.append( + ( + "LangChain + LiteLLM with Config User-Agent", + test_langchain_litellm_with_user_agent(config), + ) + ) + + # Test 3: LiteLLM Async Completion with config user agent + all_results.append( + ( + "LiteLLM Async Completion with Config User-Agent", + test_litellm_async_completion(config), + ) + ) + + # Test 4: LiteLLM Streaming Completion with config user agent + all_results.append( + ( + "LiteLLM Streaming Completion with Config User-Agent", + test_litellm_streaming_completion(config), + ) + ) + + # Test 5: LiteLLM Embedding with config user agent + all_results.append( + ( + "LiteLLM Embedding with Config User-Agent", + test_litellm_embedding_with_user_agent(config), + ) + ) + + # Summary + print("\n" + "=" * 60) + print("TEST SUMMARY") + print("=" * 60) + + passed = sum(1 for _, r in all_results if r is True) + failed = sum(1 for _, r in all_results if r is False) + skipped = sum(1 for _, r in all_results if r is None) + + for name, result in all_results: + status = ( + "✓ PASSED" + if result is True + else ("✗ FAILED" if result is False else "○ SKIPPED") + ) + print(f" {status}: {name}") + + print(f"\n Total: {passed} passed, {failed} failed, {skipped} skipped") + + if auth_method == "all": + print(f"\n Auth methods tested: {', '.join(auth_methods_to_test)}") + + return 0 if failed == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py new file mode 100644 index 00000000000..b4dc6c68bb0 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -0,0 +1,662 @@ +""" +Unit Tests for Databricks Partner Integration Features +======================================================= + +These tests are designed for automated CI/CD pipelines and do NOT require +real Databricks credentials. All external calls are mocked. + +For integration tests that use real Databricks credentials, see: + test_databricks_integration.py + +Features Tested: + - User-Agent building with partner prefixing (Databricks partner telemetry) + - Token/sensitive data redaction for secure logging + - OAuth M2M (Machine-to-Machine) authentication flow + - Databricks SDK partner telemetry registration + - Authentication priority (OAuth M2M > PAT > SDK) + +Run with: + pytest test_databricks_partner_integration.py -v + +These tests align with Databricks Partner Architecture best practices: + https://github.com/databrickslabs/partner-architecture +""" + +import json +import os +import sys + +import pytest +from unittest.mock import MagicMock, patch, Mock + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.databricks.common_utils import DatabricksBase, DatabricksException + + +class TestBuildUserAgent: + """Test cases for User-Agent string building.""" + + def test_default_user_agent(self): + """No custom user agent returns litellm/{version}.""" + ua = DatabricksBase._build_user_agent(None) + assert ua.startswith("litellm/") + assert "_" not in ua.split("/")[0] + + def test_custom_user_agent_with_version(self): + """Custom user agent with version extracts partner name.""" + ua = DatabricksBase._build_user_agent("mycompany/1.0.0") + assert ua.startswith("mycompany_litellm/") + # Verify the version is litellm's, not the custom one + assert "/1.0.0" not in ua or "mycompany_litellm/1.0.0" not in ua + + def test_custom_user_agent_without_version(self): + """Custom user agent without version still works.""" + ua = DatabricksBase._build_user_agent("mycompany") + assert ua.startswith("mycompany_litellm/") + + def test_custom_user_agent_with_underscore(self): + """Partner names with underscores are preserved.""" + ua = DatabricksBase._build_user_agent("my_company/2.0.0") + assert ua.startswith("my_company_litellm/") + + def test_custom_user_agent_with_hyphen(self): + """Partner names with hyphens are preserved.""" + ua = DatabricksBase._build_user_agent("my-company/2.0.0") + assert ua.startswith("my-company_litellm/") + + def test_custom_user_agent_ignores_custom_version(self): + """Custom version is ignored, litellm version is used.""" + ua = DatabricksBase._build_user_agent("partner/99.99.99") + parts = ua.split("/") + assert parts[0] == "partner_litellm" + assert parts[1] != "99.99.99" + + def test_empty_string_returns_default(self): + """Empty string returns default user agent.""" + ua = DatabricksBase._build_user_agent("") + assert ua.startswith("litellm/") + assert "_" not in ua.split("/")[0] + + def test_whitespace_only_returns_default(self): + """Whitespace-only string returns default user agent.""" + ua = DatabricksBase._build_user_agent(" ") + assert ua.startswith("litellm/") + assert "_" not in ua.split("/")[0] + + def test_invalid_partner_name_returns_default(self): + """Invalid partner names (special chars) return default.""" + ua = DatabricksBase._build_user_agent("my@company/1.0.0") + assert ua.startswith("litellm/") + + def test_partner_with_numbers(self): + """Partner names with numbers work.""" + ua = DatabricksBase._build_user_agent("company123/1.0.0") + assert ua.startswith("company123_litellm/") + + +class TestRedactSensitiveData: + """Test cases for sensitive data redaction.""" + + def test_redact_bearer_token_in_string(self): + """Bearer tokens are redacted in strings.""" + result = DatabricksBase.redact_sensitive_data("Bearer dapi12345abcdef") + assert "dapi12345abcdef" not in result + assert "[REDACTED]" in result + + def test_redact_dict_with_authorization(self): + """Dict with authorization key is redacted.""" + data = {"Authorization": "Bearer secret123", "other": "value"} + result = DatabricksBase.redact_sensitive_data(data) + assert result["Authorization"] == "[REDACTED]" + assert result["other"] == "value" + + def test_redact_nested_dict(self): + """Nested dicts with sensitive keys are redacted.""" + data = {"config": {"api_key": "secret", "name": "test"}} + result = DatabricksBase.redact_sensitive_data(data) + assert result["config"]["api_key"] == "[REDACTED]" + assert result["config"]["name"] == "test" + + def test_redact_pat_token(self): + """Databricks PAT tokens are redacted.""" + result = DatabricksBase.redact_sensitive_data( + "Using token dapi_fake_test_token_value" + ) + assert "dapi_fake_test_token_value" not in result + assert "[REDACTED_PAT]" in result + + def test_redact_client_secret(self): + """Client secrets are redacted.""" + data = {"client_secret": "my-super-secret-value"} + result = DatabricksBase.redact_sensitive_data(data) + assert result["client_secret"] == "[REDACTED]" + + def test_redact_list_of_dicts(self): + """Lists containing dicts with sensitive data are redacted.""" + data = [{"api_key": "secret1"}, {"name": "test"}] + result = DatabricksBase.redact_sensitive_data(data) + assert result[0]["api_key"] == "[REDACTED]" + assert result[1]["name"] == "test" + + def test_redact_none_returns_none(self): + """None input returns None.""" + assert DatabricksBase.redact_sensitive_data(None) is None + + def test_redact_preserves_non_sensitive_data(self): + """Non-sensitive data is preserved.""" + data = {"model": "dbrx", "temperature": 0.7, "messages": ["hello"]} + result = DatabricksBase.redact_sensitive_data(data) + assert result == data + + +class TestRedactHeadersForLogging: + """Test cases for header redaction.""" + + def test_authorization_header_partially_shown(self): + """Authorization header shows first 8 chars then redacts.""" + headers = {"Authorization": "Bearer dapi123456789abcdef"} + result = DatabricksBase.redact_headers_for_logging(headers) + assert result["Authorization"].startswith("Bearer d") + assert "[REDACTED]" in result["Authorization"] + + def test_short_authorization_header_fully_redacted(self): + """Short authorization values are fully redacted.""" + headers = {"Authorization": "short"} + result = DatabricksBase.redact_headers_for_logging(headers) + assert result["Authorization"] == "[REDACTED]" + + def test_non_sensitive_headers_preserved(self): + """Non-sensitive headers are not modified.""" + headers = {"Content-Type": "application/json", "User-Agent": "test/1.0"} + result = DatabricksBase.redact_headers_for_logging(headers) + assert result["Content-Type"] == "application/json" + assert result["User-Agent"] == "test/1.0" + + def test_empty_headers_returns_empty(self): + """Empty headers dict returns empty dict.""" + assert DatabricksBase.redact_headers_for_logging({}) == {} + + def test_none_headers_returns_empty(self): + """None headers returns empty dict.""" + assert DatabricksBase.redact_headers_for_logging(None) == {} + + def test_x_api_key_header_redacted(self): + """X-API-Key header is redacted.""" + headers = {"X-API-Key": "my-api-key-12345"} + result = DatabricksBase.redact_headers_for_logging(headers) + assert "[REDACTED]" in result["X-API-Key"] + + +class TestOAuthM2M: + """Test cases for OAuth M2M authentication.""" + + def test_oauth_m2m_token_success(self): + """OAuth M2M token is successfully obtained.""" + databricks_base = DatabricksBase() + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "test-access-token"} + + with patch("requests.post", return_value=mock_response) as mock_post: + token = databricks_base._get_oauth_m2m_token( + api_base="https://adb-123.azuredatabricks.net/serving-endpoints", + client_id="test-client-id", + client_secret="test-client-secret", + ) + + assert token == "test-access-token" + mock_post.assert_called_once() + call_args = mock_post.call_args + assert "oidc/v1/token" in call_args[0][0] + assert call_args[1]["data"]["grant_type"] == "client_credentials" + + def test_oauth_m2m_token_failure(self): + """OAuth M2M raises exception on failure.""" + databricks_base = DatabricksBase() + + mock_response = Mock() + mock_response.status_code = 401 + mock_response.text = "Unauthorized" + + with patch("requests.post", return_value=mock_response): + with pytest.raises(DatabricksException) as exc_info: + databricks_base._get_oauth_m2m_token( + api_base="https://adb-123.azuredatabricks.net", + client_id="bad-client-id", + client_secret="bad-secret", + ) + assert exc_info.value.status_code == 401 + + def test_oauth_m2m_strips_serving_endpoints(self): + """OAuth M2M correctly strips /serving-endpoints from URL.""" + databricks_base = DatabricksBase() + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "token"} + + with patch("requests.post", return_value=mock_response) as mock_post: + databricks_base._get_oauth_m2m_token( + api_base="https://adb-123.azuredatabricks.net/serving-endpoints", + client_id="id", + client_secret="secret", + ) + + call_url = mock_post.call_args[0][0] + assert "/serving-endpoints" not in call_url + assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token" + + +class TestValidateEnvironmentWithOAuth: + """Test OAuth M2M is used when credentials are available.""" + + def test_oauth_used_when_credentials_set(self, monkeypatch): + """OAuth M2M is used when client_id and client_secret are set.""" + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "test-client-id") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "test-secret") + monkeypatch.setenv( + "DATABRICKS_API_BASE", "https://adb-123.net/serving-endpoints" + ) + + databricks_base = DatabricksBase() + + with patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + ) as mock_oauth: + api_base, headers = databricks_base.databricks_validate_environment( + api_key=None, + api_base=None, + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + mock_oauth.assert_called_once() + assert headers["Authorization"] == "Bearer oauth-token" + + def test_pat_used_when_api_key_set(self, monkeypatch): + """PAT is used when api_key is provided.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="dapi-test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + assert headers["Authorization"] == "Bearer dapi-test-key" + + +class TestValidateEnvironmentUserAgent: + """Test User-Agent is correctly set in validate_environment.""" + + def test_default_user_agent(self, monkeypatch): + """Default user agent is set when no custom agent provided.""" + monkeypatch.delenv("DATABRICKS_USER_AGENT", raising=False) + monkeypatch.delenv("LITELLM_USER_AGENT", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + custom_user_agent=None, + ) + + assert headers["User-Agent"].startswith("litellm/") + assert "_" not in headers["User-Agent"].split("/")[0] + + def test_custom_user_agent_via_param(self, monkeypatch): + """Custom user agent is prefixed when passed as parameter.""" + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + custom_user_agent="mycompany/1.0.0", + ) + + assert headers["User-Agent"].startswith("mycompany_litellm/") + + +class TestSDKPartnerTelemetry: + """Test that SDK partner telemetry is registered.""" + + def test_sdk_partner_registered(self): + """useragent.with_partner is called when using SDK.""" + databricks_base = DatabricksBase() + + mock_workspace_client = MagicMock() + mock_workspace_client.config.host = "https://adb-123.net" + mock_workspace_client.config.authenticate.return_value = { + "Authorization": "Bearer token" + } + + with patch( + "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client + ): + with patch("databricks.sdk.useragent.with_partner") as mock_with_partner: + databricks_base._get_databricks_credentials( + api_key=None, + api_base=None, + headers=None, + ) + + mock_with_partner.assert_called_once_with("litellm") + + +class TestUserAgentFromEnvironment: + """Test User-Agent is correctly picked up from environment variables.""" + + def test_user_agent_from_databricks_env_var(self, monkeypatch): + """DATABRICKS_USER_AGENT environment variable is used.""" + monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner") + monkeypatch.delenv("LITELLM_USER_AGENT", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + custom_user_agent="envpartner", # Simulating what transformation.py passes + ) + + assert headers["User-Agent"].startswith("envpartner_litellm/") + + def test_custom_param_takes_precedence(self, monkeypatch): + """Custom user_agent parameter takes precedence over environment.""" + monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner") + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + custom_user_agent="parampartner/1.0.0", + ) + + assert headers["User-Agent"].startswith("parampartner_litellm/") + + +class TestLiteLLMCompletionUserAgent: + """Test User-Agent is correctly passed through LiteLLM completion calls.""" + + def test_completion_passes_user_agent_to_headers(self): + """litellm.completion() correctly passes user_agent to request headers.""" + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + config = DatabricksConfig() + optional_params = {"user_agent": "testpartner/1.0.0"} + + # Mock the validation to capture what headers are set + with patch.object( + config, + "databricks_validate_environment", + return_value=( + "https://test.net/serving-endpoints/chat/completions", + { + "Authorization": "Bearer test", + "User-Agent": "testpartner_litellm/1.0.0", + }, + ), + ) as mock_validate: + result = config.validate_environment( + headers={}, + model="databricks/test-model", + messages=[], + optional_params=optional_params, + litellm_params={}, + api_key="test-key", + api_base="https://test.net/serving-endpoints", + ) + + # Verify user_agent was passed to databricks_validate_environment + mock_validate.assert_called_once() + call_kwargs = mock_validate.call_args[1] + assert call_kwargs.get("custom_user_agent") == "testpartner/1.0.0" + + def test_user_agent_removed_from_optional_params(self): + """user_agent is removed from optional_params so it's not sent to API.""" + from litellm.llms.databricks.chat.transformation import DatabricksConfig + + config = DatabricksConfig() + optional_params = { + "user_agent": "testpartner/1.0.0", + "temperature": 0.7, + } + + with patch.object( + config, + "databricks_validate_environment", + return_value=( + "https://test.net/chat/completions", + {"Authorization": "Bearer test", "User-Agent": "test"}, + ), + ): + config.validate_environment( + headers={}, + model="databricks/test-model", + messages=[], + optional_params=optional_params, + litellm_params={}, + api_key="test-key", + api_base="https://test.net/serving-endpoints", + ) + + # user_agent should be removed from optional_params + assert "user_agent" not in optional_params + # Other params should remain + assert optional_params.get("temperature") == 0.7 + + +class TestLiteLLMEmbeddingUserAgent: + """Test User-Agent is correctly passed through LiteLLM embedding calls.""" + + def test_embedding_passes_user_agent_to_headers(self): + """litellm.embedding() correctly passes user_agent to request headers.""" + from litellm.llms.databricks.embed.handler import DatabricksEmbeddingHandler + + handler = DatabricksEmbeddingHandler() + optional_params = {"user_agent": "embedpartner/1.0.0"} + + with patch.object( + handler, + "databricks_validate_environment", + return_value=( + "https://test.net/serving-endpoints/embeddings", + { + "Authorization": "Bearer test", + "User-Agent": "embedpartner_litellm/1.0.0", + }, + ), + ) as mock_validate: + with patch( + "litellm.llms.openai_like.embedding.handler.OpenAILikeEmbeddingHandler.embedding" + ): + try: + handler.embedding( + model="databricks/test-model", + input=["test"], + timeout=30, + api_key="test-key", + api_base="https://test.net/serving-endpoints", + optional_params=optional_params, + ) + except Exception: + pass # We just want to verify the mock was called + + # Verify user_agent was passed + if mock_validate.called: + call_kwargs = mock_validate.call_args[1] + assert call_kwargs.get("custom_user_agent") == "embedpartner/1.0.0" + + +class TestAuthenticationPriority: + """Test that authentication methods are used in correct priority order.""" + + def test_oauth_used_when_no_api_key_provided(self, monkeypatch): + """OAuth M2M is used when OAuth creds are set and no api_key is provided.""" + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints") + + databricks_base = DatabricksBase() + + with patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + ) as mock_oauth: + api_base, headers = databricks_base.databricks_validate_environment( + api_key=None, # No PAT provided - OAuth should be used + api_base=None, + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + # OAuth should be used + mock_oauth.assert_called_once() + assert headers["Authorization"] == "Bearer oauth-token" + + def test_explicit_pat_takes_priority_over_oauth_env(self, monkeypatch): + """Explicit api_key takes priority over OAuth token in final headers.""" + monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id") + monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret") + monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints") + + databricks_base = DatabricksBase() + + # Mock the OAuth call - it will be attempted but PAT should override + with patch.object( + databricks_base, "_get_oauth_m2m_token", return_value="oauth-token" + ): + api_base, headers = databricks_base.databricks_validate_environment( + api_key="dapi-explicit-pat", + api_base=None, + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + # PAT should override OAuth token since api_key was explicitly provided + assert headers["Authorization"] == "Bearer dapi-explicit-pat" + + def test_pat_used_when_no_oauth_credentials(self, monkeypatch): + """PAT is used when OAuth credentials are not set.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="dapi-pat-token", + api_base="https://test.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + assert headers["Authorization"] == "Bearer dapi-pat-token" + + def test_sdk_fallback_when_no_credentials(self, monkeypatch): + """Databricks SDK is used when no API key or OAuth credentials.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + monkeypatch.delenv("DATABRICKS_API_KEY", raising=False) + + databricks_base = DatabricksBase() + + mock_workspace_client = MagicMock() + mock_workspace_client.config.host = "https://adb-123.net" + mock_workspace_client.config.authenticate.return_value = { + "Authorization": "Bearer sdk-token" + } + + with patch( + "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client + ): + with patch("databricks.sdk.useragent.with_partner"): + api_base, headers = databricks_base.databricks_validate_environment( + api_key=None, + api_base=None, + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + assert "Authorization" in headers + + +class TestEndpointURLConstruction: + """Test that endpoint URLs are correctly constructed.""" + + def test_chat_completions_endpoint(self, monkeypatch): + """Chat completions endpoint is correctly appended.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://test.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) + + assert api_base.endswith("/chat/completions") + + def test_embeddings_endpoint(self, monkeypatch): + """Embeddings endpoint is correctly appended.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://test.net/serving-endpoints", + endpoint_type="embeddings", + custom_endpoint=False, + headers=None, + ) + + assert api_base.endswith("/embeddings") + + def test_custom_endpoint_not_modified(self, monkeypatch): + """Custom endpoints are not modified.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://test.net/custom/endpoint", + endpoint_type="chat_completions", + custom_endpoint=True, + headers=None, + ) + + assert api_base == "https://test.net/custom/endpoint" From 87fc81f3e6f21b116fad941dfa9ee7ab91cf1f5b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 13:19:03 -0800 Subject: [PATCH 025/330] Add cloudzero ui docs --- .../docs/observability/cloudzero.md | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/docs/my-website/docs/observability/cloudzero.md b/docs/my-website/docs/observability/cloudzero.md index f213ef64e13..19f6d80ca8b 100644 --- a/docs/my-website/docs/observability/cloudzero.md +++ b/docs/my-website/docs/observability/cloudzero.md @@ -65,6 +65,52 @@ Start your LiteLLM proxy with the configuration: litellm --config /path/to/config.yaml ``` +## Setup on UI + +1\. Click "Settings" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/5ac36280-c688-41a3-8d0e-23e19c6a470b/ascreenshot.jpeg?tl_px=0,332&br_px=1308,1064&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=119,444) + + +2\. Click "Logging & Alerts" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/13f76b09-e0c4-4738-ba05-2d5111c6ad3e/ascreenshot.jpeg?tl_px=0,332&br_px=1308,1064&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=58,507) + + +3\. Click "CloudZero Cost Tracking" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/f96cc1e5-7bc0-4d7c-9aeb-5cbbec549b12/ascreenshot.jpeg?tl_px=0,0&br_px=1308,731&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=389,56) + + +4\. Click "Add CloudZero Integration" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/04fbc748-0e6f-43bb-8a57-dd2e83dbfcb5/ascreenshot.jpeg?tl_px=0,90&br_px=1308,821&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=616,277) + + +5\. Enter your CloudZero API Key. + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/080e82f1-f94f-4ed7-8014-e495380336f3/ascreenshot.jpeg?tl_px=0,0&br_px=1308,731&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=506,129) + + +6\. Enter your CloudZero Connection ID. + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/af417aa2-67a8-4dee-a014-84b1892dc07e/ascreenshot.jpeg?tl_px=0,0&br_px=1308,731&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=488,213) + + +7\. Click "Create" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/647e672f-9a4a-4754-a7b0-abf1397abad4/ascreenshot.jpeg?tl_px=0,88&br_px=1308,819&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=711,277) + + +8\. Test your payload with "Run Dry Run Simulation" + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/7447cbe0-3450-4be5-bdc4-37fb8280aa58/ascreenshot.jpeg?tl_px=0,125&br_px=1308,856&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=334,277) + + +10\. Click "Export Data Now" to export to CLoudZero + +![](https://ajeuwbhvhr.cloudimg.io/https://colony-recorder.s3.amazonaws.com/files/2025-12-22/7be9bd48-6e27-4c68-bc75-946f3ab593d9/ascreenshot.jpeg?tl_px=0,130&br_px=1308,861&force_format=jpeg&q=100&width=1120.0&wat=1&wat_opacity=0.7&wat_gravity=northwest&wat_url=https://colony-recorder.s3.us-west-1.amazonaws.com/images/watermarks/FB923C_standard.png&wat_pad=518,277) + ## Testing Your Setup ### Dry Run Export From dc3bdffaee2d648d367523f780a06557a8022f52 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 14:33:38 -0800 Subject: [PATCH 026/330] Migrate MCP servers to react query --- .../hooks/mcpServers/useMCPAccessGroups.ts | 13 +++ .../hooks/mcpServers/useMCPServers.ts | 14 +++ .../MCPServerSelector.tsx | 41 ++------- .../MCPToolPermissions.test.tsx | 43 +++++---- .../MCPToolPermissions.tsx | 73 +++++---------- .../src/components/mcp_tools/mcp_servers.tsx | 18 +--- .../organisms/create_key_button.tsx | 92 +++++++------------ .../src/components/team/team_info.test.tsx | 15 +-- .../templates/key_edit_view.test.tsx | 9 +- ui/litellm-dashboard/tests/test-utils.tsx | 21 ++++- 10 files changed, 155 insertions(+), 184 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts new file mode 100644 index 00000000000..eeeb76bb742 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups.ts @@ -0,0 +1,13 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPAccessGroups } from "@/components/networking"; + +const mcpAccessGroupsKeys = createQueryKeys("mcpAccessGroups"); + +export const useMCPAccessGroups = (accessToken: string | null) => { + return useQuery({ + queryKey: mcpAccessGroupsKeys.list({}), + queryFn: async () => await fetchMCPAccessGroups(accessToken!), + enabled: !!accessToken, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts new file mode 100644 index 00000000000..02e471d8e5f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServers.ts @@ -0,0 +1,14 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPServers } from "@/components/networking"; +import { MCPServer } from "@/components/mcp_tools/types"; + +const mcpServersKeys = createQueryKeys("mcpServers"); + +export const useMCPServers = (accessToken: string | null) => { + return useQuery({ + queryKey: mcpServersKeys.list({}), + queryFn: async () => await fetchMCPServers(accessToken!), + enabled: !!accessToken, + }); +}; diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index 7b79f5cc707..f30a76a229b 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -1,15 +1,13 @@ -import React, { useEffect, useState } from "react"; +import React from "react"; import { Select } from "antd"; -import { fetchMCPServers, fetchMCPAccessGroups } from "../networking"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPAccessGroups } from "@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups"; import { MCPServer } from "../mcp_tools/types"; interface MCPServerSelectorProps { - onChange: (selected: { - servers: string[]; - accessGroups: string[]; - }) => void; - value?: { - servers: string[]; + onChange: (selected: { servers: string[]; accessGroups: string[] }) => void; + value?: { + servers: string[]; accessGroups: string[]; }; className?: string; @@ -26,31 +24,10 @@ const MCPServerSelector: React.FC = ({ placeholder = "Select MCP servers", disabled = false, }) => { - const [mcpServers, setMCPServers] = useState([]); - const [accessGroups, setAccessGroups] = useState([]); - const [loading, setLoading] = useState(false); + const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(accessToken); + const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(accessToken); - useEffect(() => { - const fetchData = async () => { - if (!accessToken) return; - setLoading(true); - try { - const [serversRes, groupsRes] = await Promise.all([ - fetchMCPServers(accessToken), - fetchMCPAccessGroups(accessToken), - ]); - let servers = Array.isArray(serversRes) ? serversRes : serversRes.data || []; - let groups = Array.isArray(groupsRes) ? groupsRes : groupsRes.data || []; - setMCPServers(servers); - setAccessGroups(groups); - } catch (error) { - console.error("Error fetching MCP servers or access groups:", error); - } finally { - setLoading(false); - } - }; - fetchData(); - }, [accessToken]); + const loading = serversLoading || groupsLoading; // Combine options, access groups first const options = [ diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index 93f96966d0a..ec541d7a63c 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -1,11 +1,22 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import MCPToolPermissions from "./MCPToolPermissions"; import * as networking from "../networking"; vi.mock("../networking"); +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + describe("MCPToolPermissions", () => { const mockAccessToken = "test-token"; const mockServerId = "server-123"; @@ -28,15 +39,13 @@ describe("MCPToolPermissions", () => { ]; // Mock fetchMCPServers to return server details - vi.mocked(networking.fetchMCPServers).mockResolvedValue({ - data: [ - { - server_id: mockServerId, - server_name: mockServerName, - alias: mockServerName, - }, - ], - }); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + server_id: mockServerId, + server_name: mockServerName, + alias: mockServerName, + }, + ]); // Mock listMCPTools to return tools for the server vi.mocked(networking.listMCPTools).mockResolvedValue({ @@ -44,13 +53,16 @@ describe("MCPToolPermissions", () => { error: false, }); + const queryClient = createQueryClient(); render( - + + + , ); // Wait for server and tools to load @@ -76,4 +88,3 @@ describe("MCPToolPermissions", () => { expect(networking.listMCPTools).toHaveBeenCalledWith(mockAccessToken, mockServerId); }); }); - diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index 3f36caf4670..a567791e221 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -1,9 +1,10 @@ -import React, { useEffect, useState } from "react"; -import { listMCPTools, fetchMCPServers } from "../networking"; +import React, { useEffect, useState, useMemo } from "react"; +import { listMCPTools } from "../networking"; import { MCPTool, MCPServer } from "../mcp_tools/types"; import { Text } from "@tremor/react"; import { Spin, Checkbox } from "antd"; import { XIcon } from "lucide-react"; +import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; interface MCPToolPermissionsProps { accessToken: string; @@ -20,63 +21,43 @@ const MCPToolPermissions: React.FC = ({ onChange, disabled = false, }) => { - const [servers, setServers] = useState([]); + const { data: allServers = [] } = useMCPServers(accessToken); const [serverTools, setServerTools] = useState>({}); const [loadingTools, setLoadingTools] = useState>({}); const [toolErrors, setToolErrors] = useState>({}); - // Fetch server details - useEffect(() => { - const loadServerDetails = async () => { - if (selectedServers.length === 0) { - setServers([]); - return; - } - - try { - const response = await fetchMCPServers(accessToken); - const allServers = Array.isArray(response) ? response : response.data || []; - - const filteredServers = allServers.filter((server: MCPServer) => - selectedServers.includes(server.server_id) - ); - - setServers(filteredServers); - } catch (error) { - console.error("Error fetching MCP servers:", error); - setServers([]); - } - }; - - loadServerDetails(); - }, [selectedServers, accessToken]); + // Filter servers based on selectedServers + const servers = useMemo(() => { + if (selectedServers.length === 0) return []; + return allServers.filter((server: MCPServer) => selectedServers.includes(server.server_id)); + }, [allServers, selectedServers]); // Fetch tools for a specific server const fetchToolsForServer = async (serverId: string) => { - setLoadingTools(prev => ({ ...prev, [serverId]: true })); - setToolErrors(prev => ({ ...prev, [serverId]: "" })); - + setLoadingTools((prev) => ({ ...prev, [serverId]: true })); + setToolErrors((prev) => ({ ...prev, [serverId]: "" })); + try { const response = await listMCPTools(accessToken, serverId); - + if (response.error) { - setToolErrors(prev => ({ ...prev, [serverId]: response.message || "Failed to fetch tools" })); - setServerTools(prev => ({ ...prev, [serverId]: [] })); + setToolErrors((prev) => ({ ...prev, [serverId]: response.message || "Failed to fetch tools" })); + setServerTools((prev) => ({ ...prev, [serverId]: [] })); } else { - setServerTools(prev => ({ ...prev, [serverId]: response.tools || [] })); + setServerTools((prev) => ({ ...prev, [serverId]: response.tools || [] })); } } catch (err) { console.error(`Error fetching tools for server ${serverId}:`, err); - setToolErrors(prev => ({ ...prev, [serverId]: "Failed to fetch tools" })); - setServerTools(prev => ({ ...prev, [serverId]: [] })); + setToolErrors((prev) => ({ ...prev, [serverId]: "Failed to fetch tools" })); + setServerTools((prev) => ({ ...prev, [serverId]: [] })); } finally { - setLoadingTools(prev => ({ ...prev, [serverId]: false })); + setLoadingTools((prev) => ({ ...prev, [serverId]: false })); } }; // Auto-fetch tools when servers change useEffect(() => { - servers.forEach(server => { + servers.forEach((server) => { if (!serverTools[server.server_id] && !loadingTools[server.server_id]) { fetchToolsForServer(server.server_id); } @@ -87,9 +68,9 @@ const MCPToolPermissions: React.FC = ({ const handleToolToggle = (serverId: string, toolName: string) => { const currentTools = toolPermissions[serverId] || []; const newTools = currentTools.includes(toolName) - ? currentTools.filter(name => name !== toolName) + ? currentTools.filter((name) => name !== toolName) : [...currentTools, toolName]; - + const updatedPermissions = { ...toolPermissions, [serverId]: newTools, @@ -101,7 +82,7 @@ const MCPToolPermissions: React.FC = ({ const tools = serverTools[serverId] || []; onChange({ ...toolPermissions, - [serverId]: tools.map(t => t.name), + [serverId]: tools.map((t) => t.name), }); }; @@ -131,9 +112,7 @@ const MCPToolPermissions: React.FC = ({
{serverName} - {server.description && ( - {server.description} - )} + {server.description && {server.description}}
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index a46738ab365..83393c4a94b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -1,11 +1,11 @@ import { isAdminRole } from "@/utils/roles"; import { QuestionCircleOutlined } from "@ant-design/icons"; -import { useQuery } from "@tanstack/react-query"; import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; import { Descriptions, Modal, Select, Tooltip, Typography } from "antd"; import React, { useEffect, useState } from "react"; +import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers"; import NotificationsManager from "../molecules/notifications_manager"; -import { deleteMCPServer, fetchMCPServers } from "../networking"; +import { deleteMCPServer } from "../networking"; import { DataTable } from "../view_logs/table"; import CreateMCPServer from "./create_mcp_server"; import MCPConnect from "./mcp_connect"; @@ -19,19 +19,7 @@ const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { - const { - data: mcpServers, - isLoading: isLoadingServers, - refetch, - dataUpdatedAt, - } = useQuery({ - queryKey: ["mcpServers"], - queryFn: () => { - if (!accessToken) throw new Error("Access Token required"); - return fetchMCPServers(accessToken); - }, - enabled: !!accessToken, - }) as { data: MCPServer[]; isLoading: boolean; refetch: () => void; dataUpdatedAt: number }; + const { data: mcpServers, isLoading: isLoadingServers, refetch, dataUpdatedAt } = useMCPServers(accessToken); // Log allowed_tools from fetched servers React.useEffect(() => { diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 3c0a0f520d4..f3ad803e27c 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1,43 +1,40 @@ "use client"; -import React, { useState, useEffect, useCallback } from "react"; -import { Button, TextInput, Grid, Col } from "@tremor/react"; -import { Text, Title, Accordion, AccordionHeader, AccordionBody } from "@tremor/react"; -import { CopyToClipboard } from "react-copy-to-clipboard"; -import { Button as Button2, Modal, Form, Input, Select, Radio, Switch } from "antd"; -import NumericalInput from "../shared/numerical_input"; -import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; -import SchemaFormFields from "../common_components/check_openapi_schema"; -import { - keyCreateCall, - modelAvailableCall, - getGuardrailsList, - proxyBaseUrl, - getPossibleUserRoles, - userFilterUICall, - keyCreateServiceAccountCall, - fetchMCPAccessGroups, - getPromptsList, -} from "../networking"; -import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; -import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; -import { Team } from "../key_team_helpers/key_list"; -import TeamDropdown from "../common_components/team_dropdown"; -import { InfoCircleOutlined } from "@ant-design/icons"; -import { Tooltip } from "antd"; -import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"; -import Createuser from "../create_user_button"; -import debounce from "lodash/debounce"; -import { rolesWithWriteAccess } from "../../utils/roles"; -import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react"; +import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tooltip } from "antd"; +import debounce from "lodash/debounce"; +import React, { useCallback, useEffect, useState } from "react"; +import { CopyToClipboard } from "react-copy-to-clipboard"; +import { rolesWithWriteAccess } from "../../utils/roles"; +import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; +import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; +import SchemaFormFields from "../common_components/check_openapi_schema"; +import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; +import ModelAliasManager from "../common_components/ModelAliasManager"; +import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; +import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"; +import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; +import TeamDropdown from "../common_components/team_dropdown"; +import Createuser from "../create_user_button"; +import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { Team } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; -import AgentSelector from "../agent_management/AgentSelector"; -import ModelAliasManager from "../common_components/ModelAliasManager"; import NotificationsManager from "../molecules/notifications_manager"; -import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; -import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; +import { + getGuardrailsList, + getPossibleUserRoles, + getPromptsList, + keyCreateCall, + keyCreateServiceAccountCall, + modelAvailableCall, + proxyBaseUrl, + userFilterUICall, +} from "../networking"; +import NumericalInput from "../shared/numerical_input"; +import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; const { Option } = Select; @@ -168,7 +165,6 @@ const CreateKey: React.FC = ({ const [userOptions, setUserOptions] = useState([]); const [userSearchLoading, setUserSearchLoading] = useState(false); const [mcpAccessGroups, setMcpAccessGroups] = useState([]); - const [mcpAccessGroupsLoaded, setMcpAccessGroupsLoaded] = useState(false); const [disabledCallbacks, setDisabledCallbacks] = useState([]); const [keyType, setKeyType] = useState("default"); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); @@ -205,22 +201,6 @@ const CreateKey: React.FC = ({ } }, [accessToken, userID, userRole]); - const fetchMcpAccessGroups = async () => { - try { - if (accessToken == null) { - return; - } - const groups = await fetchMCPAccessGroups(accessToken); - setMcpAccessGroups(groups); - } catch (error) { - console.error("Failed to fetch MCP access groups:", error); - } - }; - - useEffect(() => { - fetchMcpAccessGroups(); - }, [accessToken]); - useEffect(() => { const fetchGuardrails = async () => { try { @@ -1053,15 +1033,7 @@ const CreateKey: React.FC = ({ options={predefinedTags} /> - { - if (!mcpAccessGroupsLoaded) { - fetchMcpAccessGroups(); - setMcpAccessGroupsLoaded(true); - } - }} - > + MCP Settings diff --git a/ui/litellm-dashboard/src/components/team/team_info.test.tsx b/ui/litellm-dashboard/src/components/team/team_info.test.tsx index 9b19611828f..4f7d70ba6ab 100644 --- a/ui/litellm-dashboard/src/components/team/team_info.test.tsx +++ b/ui/litellm-dashboard/src/components/team/team_info.test.tsx @@ -1,5 +1,6 @@ import * as networking from "@/components/networking"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import { renderWithProviders } from "../../../tests/test-utils"; import { afterEach, describe, expect, it, vi } from "vitest"; import TeamInfoView from "./team_info"; @@ -62,7 +63,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); - render( + renderWithProviders( {}} @@ -124,7 +125,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); - render( + renderWithProviders( {}} @@ -219,7 +220,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); - render( + renderWithProviders( {}} @@ -310,7 +311,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); - render( + renderWithProviders( {}} @@ -373,7 +374,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: teamResponse.team_info, team_id: "123" } as any); - render( + renderWithProviders( {}} @@ -450,7 +451,7 @@ describe("TeamInfoView", () => { vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: teamResponse.team_info, team_id: "123" } as any); - render( + renderWithProviders( {}} diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index e23bc88dda8..85c6192693e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1,4 +1,5 @@ -import { render, waitFor } from "@testing-library/react"; +import { waitFor } from "@testing-library/react"; +import { renderWithProviders } from "../../../tests/test-utils"; import { describe, expect, it, vi } from "vitest"; import { KeyEditView } from "./key_edit_view"; import { KeyResponse } from "../key_team_helpers/key_list"; @@ -89,7 +90,7 @@ describe("KeyEditView", () => { key_rotation_at: undefined, }; it("should render", async () => { - const { getByText } = render( + const { getByText } = renderWithProviders( {}} @@ -107,7 +108,7 @@ describe("KeyEditView", () => { }); it("should render tags", async () => { - const { getByText } = render( + const { getByText } = renderWithProviders( {}} @@ -125,7 +126,7 @@ describe("KeyEditView", () => { }); it("should not render tags in metadata textarea", async () => { - const { getByLabelText } = render( + const { getByLabelText } = renderWithProviders( {}} diff --git a/ui/litellm-dashboard/tests/test-utils.tsx b/ui/litellm-dashboard/tests/test-utils.tsx index cf7fbaf0d8f..ed1f248648e 100644 --- a/ui/litellm-dashboard/tests/test-utils.tsx +++ b/ui/litellm-dashboard/tests/test-utils.tsx @@ -1,9 +1,26 @@ import React, { PropsWithChildren } from "react"; import { render, RenderOptions } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; + +// Create a client for testing +const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: Infinity, + staleTime: Infinity, + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + }, + mutations: { + retry: false, + }, + }, +}); const Providers: React.FC = ({ children }) => { - // Add future providers here (Theme/Router/QueryClient/etc.) - return <>{children}; + return {children}; }; export const renderWithProviders = (ui: React.ReactElement, options?: RenderOptions) => From f258dbb03a6d967ba670944a352b7c87df86c82b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 14:40:15 -0800 Subject: [PATCH 027/330] Fixing build --- .../components/mcp_server_management/MCPServerSelector.tsx | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index f30a76a229b..7830edf5867 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -1,8 +1,7 @@ -import React from "react"; -import { Select } from "antd"; -import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPAccessGroups } from "@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups"; -import { MCPServer } from "../mcp_tools/types"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { Select } from "antd"; +import React from "react"; interface MCPServerSelectorProps { onChange: (selected: { servers: string[]; accessGroups: string[] }) => void; From 78e3ae7bdeb115aa425de9481157fe0d5b123037 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 22 Dec 2025 15:17:36 -0800 Subject: [PATCH 028/330] Fix MCP Select button submitting form --- .../MCPToolPermissions.test.tsx | 133 +++++++++++++++--- .../MCPToolPermissions.tsx | 13 +- 2 files changed, 120 insertions(+), 26 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index ec541d7a63c..fdd69d064d3 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -1,22 +1,12 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, waitFor } from "@testing-library/react"; +import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderWithProviders } from "../../../tests/test-utils"; import MCPToolPermissions from "./MCPToolPermissions"; import * as networking from "../networking"; vi.mock("../networking"); -const createQueryClient = () => - new QueryClient({ - defaultOptions: { - queries: { - retry: false, - gcTime: 0, - }, - }, - }); - describe("MCPToolPermissions", () => { const mockAccessToken = "test-token"; const mockServerId = "server-123"; @@ -53,16 +43,13 @@ describe("MCPToolPermissions", () => { error: false, }); - const queryClient = createQueryClient(); - render( - - - , + renderWithProviders( + , ); // Wait for server and tools to load @@ -87,4 +74,106 @@ describe("MCPToolPermissions", () => { expect(networking.fetchMCPServers).toHaveBeenCalledWith(mockAccessToken); expect(networking.listMCPTools).toHaveBeenCalledWith(mockAccessToken, mockServerId); }); + + it("should select all tools when Select All button is clicked", async () => { + const mockOnChange = vi.fn(); + const mockTools = [ + { name: "read_wiki_structure", description: "Get documentation topics" }, + { name: "read_wiki_contents", description: "View documentation" }, + { name: "ask_question", description: "Ask questions" }, + ]; + + // Mock fetchMCPServers to return server details + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + server_id: mockServerId, + server_name: mockServerName, + alias: mockServerName, + }, + ]); + + // Mock listMCPTools to return tools for the server + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: mockTools, + error: false, + }); + + renderWithProviders( + , + ); + + // Wait for server and tools to load + await waitFor(() => { + expect(screen.getByText(mockServerName)).toBeInTheDocument(); + }); + + await waitFor(() => { + expect(screen.getByText("read_wiki_structure")).toBeInTheDocument(); + }); + + // Click the Select All button + const selectAllButton = screen.getByRole("button", { name: "Select All" }); + await userEvent.click(selectAllButton); + + // Verify onChange was called with all tools selected + expect(mockOnChange).toHaveBeenCalledWith({ + [mockServerId]: ["read_wiki_structure", "read_wiki_contents", "ask_question"], + }); + }); + + it("should deselect all tools when Deselect All button is clicked", async () => { + const mockOnChange = vi.fn(); + const mockTools = [ + { name: "read_wiki_structure", description: "Get documentation topics" }, + { name: "read_wiki_contents", description: "View documentation" }, + { name: "ask_question", description: "Ask questions" }, + ]; + + // Mock fetchMCPServers to return server details + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + server_id: mockServerId, + server_name: mockServerName, + alias: mockServerName, + }, + ]); + + // Mock listMCPTools to return tools for the server + vi.mocked(networking.listMCPTools).mockResolvedValue({ + tools: mockTools, + error: false, + }); + + renderWithProviders( + , + ); + + // Wait for server and tools to load + await waitFor(() => { + expect(screen.getByText(mockServerName)).toBeInTheDocument(); + }); + + await waitFor(() => { + expect(screen.getByText("read_wiki_structure")).toBeInTheDocument(); + }); + + // Click the Deselect All button + const deselectAllButton = screen.getByRole("button", { name: "Deselect All" }); + await userEvent.click(deselectAllButton); + + // Verify onChange was called with no tools selected + expect(mockOnChange).toHaveBeenCalledWith({ + [mockServerId]: [], + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index a567791e221..ec7e2797814 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -80,17 +80,19 @@ const MCPToolPermissions: React.FC = ({ const handleSelectAll = (serverId: string) => { const tools = serverTools[serverId] || []; - onChange({ + const newPermissions = { ...toolPermissions, [serverId]: tools.map((t) => t.name), - }); + }; + onChange(newPermissions); }; const handleDeselectAll = (serverId: string) => { - onChange({ + const newPermissions = { ...toolPermissions, [serverId]: [], - }); + }; + onChange(newPermissions); }; if (selectedServers.length === 0) { @@ -116,6 +118,7 @@ const MCPToolPermissions: React.FC = ({
diff --git a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx index 344ff2e94f2..79224edba43 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx @@ -3,10 +3,13 @@ import { flexRender, getCoreRowModel, getSortedRowModel, + getPaginationRowModel, SortingState, useReactTable, ColumnResizeMode, VisibilityState, + PaginationState, + OnChangeFn, } from "@tanstack/react-table"; import React from "react"; import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; @@ -23,16 +26,20 @@ interface ModelDataTableProps { data: TData[]; columns: ColumnDef[]; isLoading?: boolean; - table: any; // Add table prop to access column visibility controls defaultSorting?: SortingState; + pagination?: PaginationState; + onPaginationChange?: OnChangeFn; + enablePagination?: boolean; } export function ModelDataTable({ data = [], columns, isLoading = false, - table, defaultSorting = [], + pagination, + onPaginationChange, + enablePagination = false, }: ModelDataTableProps) { const [sorting, setSorting] = React.useState(defaultSorting); const [columnResizeMode] = React.useState("onChange"); @@ -46,13 +53,16 @@ export function ModelDataTable({ sorting, columnSizing, columnVisibility, + ...(enablePagination && pagination ? { pagination } : {}), }, columnResizeMode, onSortingChange: setSorting, onColumnSizingChange: setColumnSizing, onColumnVisibilityChange: setColumnVisibility, + ...(enablePagination && onPaginationChange ? { onPaginationChange } : {}), getCoreRowModel: getCoreRowModel(), getSortedRowModel: getSortedRowModel(), + ...(enablePagination ? { getPaginationRowModel: getPaginationRowModel() } : {}), enableSorting: true, enableColumnResizing: true, defaultColumn: { @@ -61,13 +71,6 @@ export function ModelDataTable({ }, }); - // Expose table instance to parent - React.useEffect(() => { - if (table) { - table.current = tableInstance; - } - }, [tableInstance, table]); - const getHeaderText = (header: any): string => { if (typeof header === "string") { return header; From 5d1fe86cda3609071b29b4b3900cc58002184ff1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 24 Dec 2025 14:33:27 -0800 Subject: [PATCH 085/330] Tests --- .../components/AllModelsTab.test.tsx | 103 ++++++++++++------ 1 file changed, 72 insertions(+), 31 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index a4bb20128e0..dfa400e6ea9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -1,9 +1,27 @@ import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized"; -import * as useTeamsModule from "@/app/(dashboard)/hooks/useTeams"; import { render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AllModelsTab from "./AllModelsTab"; +// Mock the useModelsInfo hook +const mockUseModelsInfo = vi.fn(() => ({ data: { data: [] } })) as any; + +vi.mock("../../hooks/models/useModels", () => ({ + useModelsInfo: () => mockUseModelsInfo(), +})); + +// Mock the useTeams hook (react-query implementation) +const mockUseTeams = vi.fn(() => ({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), +})) as any; + +vi.mock("../../hooks/teams/useTeams", () => ({ + useTeams: () => mockUseTeams(), +})); + describe("AllModelsTab", () => { const mockSetSelectedModelGroup = vi.fn(); const mockSetSelectedModelId = vi.fn(); @@ -18,9 +36,6 @@ describe("AllModelsTab", () => { setSelectedModelId: mockSetSelectedModelId, setSelectedTeamId: mockSetSelectedTeamId, setEditModel: mockSetEditModel, - modelData: { - data: [], - }, }; const mockUseAuthorized = { @@ -40,9 +55,13 @@ describe("AllModelsTab", () => { }); it("should render with empty data", () => { - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: [], - setTeams: vi.fn(), + mockUseModelsInfo.mockReturnValueOnce({ data: { data: [] } }); + + mockUseTeams.mockReturnValueOnce({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), }); render(); @@ -66,9 +85,11 @@ describe("AllModelsTab", () => { }, ]; - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: mockTeams, - setTeams: vi.fn(), + mockUseTeams.mockReturnValueOnce({ + data: mockTeams, + isLoading: false, + error: null, + refetch: vi.fn(), }); const modelData = { @@ -92,7 +113,9 @@ describe("AllModelsTab", () => { ], }; - render(); + mockUseModelsInfo.mockReturnValue({ data: modelData }); + + render(); await waitFor(() => { expect(screen.getByText("Showing 0 results")).toBeInTheDocument(); @@ -116,9 +139,11 @@ describe("AllModelsTab", () => { }, ]; - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: mockTeams, - setTeams: vi.fn(), + mockUseTeams.mockReturnValue({ + data: mockTeams, + isLoading: false, + error: null, + refetch: vi.fn(), }); const modelData = { @@ -142,7 +167,9 @@ describe("AllModelsTab", () => { ], }; - render(); + mockUseModelsInfo.mockReturnValue({ data: modelData }); + + render(); await waitFor(() => { expect(screen.getByText("Showing 0 results")).toBeInTheDocument(); @@ -150,9 +177,11 @@ describe("AllModelsTab", () => { }); it("should filter models by direct_access for personal team", async () => { - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: [], - setTeams: vi.fn(), + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), }); const modelData = { @@ -178,7 +207,9 @@ describe("AllModelsTab", () => { ], }; - render(); + mockUseModelsInfo.mockReturnValue({ data: modelData }); + + render(); await waitFor(() => { expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument(); @@ -186,9 +217,11 @@ describe("AllModelsTab", () => { }); it("should show config model status for models defined in configs", async () => { - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: [], - setTeams: vi.fn(), + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), }); const modelData = { @@ -226,7 +259,9 @@ describe("AllModelsTab", () => { ], }; - render(); + mockUseModelsInfo.mockReturnValue({ data: modelData }); + + render(); await waitFor(() => { expect(screen.getByText("Config Model")).toBeInTheDocument(); @@ -235,19 +270,21 @@ describe("AllModelsTab", () => { }); it("should show 'Defined in config' for models defined in configs", async () => { - vi.spyOn(useTeamsModule, "default").mockReturnValue({ - teams: [], - setTeams: vi.fn(), + mockUseTeams.mockReturnValue({ + data: [], + isLoading: false, + error: null, + refetch: vi.fn(), }); const modelData = { data: [ { - model_name: "gpt-4-config-model", - litellm_model_name: "gpt-4-config-model", + model_name: "gpt-4-config", + litellm_model_name: "gpt-4-config", provider: "openai", model_info: { - id: "model-config-defined", + id: "model-config-1", db_model: false, direct_access: true, access_via_team_ids: [], @@ -260,8 +297,12 @@ describe("AllModelsTab", () => { ], }; - render(); + mockUseModelsInfo.mockReturnValue({ data: modelData }); - expect(screen.getByText("Defined in config")).toBeInTheDocument(); + render(); + + await waitFor(() => { + expect(screen.getByText("Defined in config")).toBeInTheDocument(); + }); }); }); From 4ce135727f8b00e6fa1af077be15b2e2accc5983 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 24 Dec 2025 14:51:00 -0800 Subject: [PATCH 086/330] Fixing build --- .../model_dashboard/HealthCheckComponent.tsx | 1 - .../src/components/model_hub_table.tsx | 9 +- .../src/components/public_model_hub.tsx | 37 +++----- .../components/templates/model_dashboard.tsx | 93 ++++++++++--------- 4 files changed, 66 insertions(+), 74 deletions(-) diff --git a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx index 994ea8adfc0..5d35b92684c 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx @@ -596,7 +596,6 @@ const HealthCheckComponent: React.FC = ({ }; })} isLoading={false} - table={healthTableRef} /> diff --git a/ui/litellm-dashboard/src/components/model_hub_table.tsx b/ui/litellm-dashboard/src/components/model_hub_table.tsx index 7d48bf68aed..f45e44ce905 100644 --- a/ui/litellm-dashboard/src/components/model_hub_table.tsx +++ b/ui/litellm-dashboard/src/components/model_hub_table.tsx @@ -1,10 +1,9 @@ import { CopyOutlined } from "@ant-design/icons"; -import { Table as TableInstance } from "@tanstack/react-table"; import { Badge, Button, Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; import { Modal } from "antd"; import { Copy } from "lucide-react"; import { useRouter } from "next/navigation"; -import React, { useCallback, useEffect, useRef, useState } from "react"; +import React, { useCallback, useEffect, useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { isAdminRole } from "../utils/roles"; import { agentHubColumns, AgentHubData } from "./agent_hub_table_columns"; @@ -76,9 +75,6 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, const [isMcpModalVisible, setIsMcpModalVisible] = useState(false); const [isMakeMcpPublicModalVisible, setIsMakeMcpPublicModalVisible] = useState(false); const router = useRouter(); - const tableRef = useRef>(null); - const agentTableRef = useRef>(null); - const mcpTableRef = useRef>(null); useEffect(() => { const fetchData = async (accessToken: string) => { @@ -404,7 +400,6 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, columns={modelHubColumns(showModal, copyToClipboard, publicPage)} data={filteredData} isLoading={loading} - table={tableRef} defaultSorting={[{ id: "model_group", desc: false }]} /> @@ -431,7 +426,6 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, columns={agentHubColumns(showAgentModal, copyToClipboard, publicPage)} data={agentHubData || []} isLoading={agentLoading} - table={agentTableRef} defaultSorting={[{ id: "name", desc: false }]} /> @@ -458,7 +452,6 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, columns={mcpHubColumns(showMcpModal, copyToClipboard, publicPage)} data={mcpHubData || []} isLoading={mcpLoading} - table={mcpTableRef} defaultSorting={[{ id: "server_name", desc: false }]} /> diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 3493f0bf93f..4678dbe3f94 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -1,25 +1,24 @@ -import React, { useEffect, useState, useRef, useMemo } from "react"; -import { - modelHubPublicModelsCall, - getPublicModelHubInfo, - agentHubPublicModelsCall, - mcpHubPublicServersCall, - getUiConfig, -} from "./networking"; -import { ModelDataTable } from "./model_dashboard/table"; -import { ColumnDef } from "@tanstack/react-table"; -import { Card, Text, Title, Button } from "@tremor/react"; -import { Tag, Tooltip, Modal, Select, Tabs } from "antd"; +import { ThemeProvider } from "@/contexts/ThemeContext"; import { ExternalLinkIcon, SearchIcon } from "@heroicons/react/outline"; +import { ColumnDef } from "@tanstack/react-table"; +import { Button, Card, Text, Title } from "@tremor/react"; +import { Modal, Select, Tabs, Tag, Tooltip } from "antd"; import { Copy, Info } from "lucide-react"; -import { Table as TableInstance } from "@tanstack/react-table"; +import React, { useEffect, useMemo, useState } from "react"; +import { ModelDataTable } from "./model_dashboard/table"; +import NotificationsManager from "./molecules/notifications_manager"; +import Navbar from "./navbar"; +import { + agentHubPublicModelsCall, + getPublicModelHubInfo, + getUiConfig, + mcpHubPublicServersCall, + modelHubPublicModelsCall, +} from "./networking"; import { generateCodeSnippet } from "./playground/chat_ui/CodeSnippets"; import { getEndpointType } from "./playground/chat_ui/mode_endpoint_mapping"; import { MessageType } from "./playground/chat_ui/types"; import { getProviderLogoAndName } from "./provider_info_helpers"; -import Navbar from "./navbar"; -import { ThemeProvider } from "@/contexts/ThemeContext"; -import NotificationsManager from "./molecules/notifications_manager"; const { TabPane } = Tabs; @@ -118,9 +117,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const [selectedMcpServer, setSelectedMcpServer] = useState(null); const [proxySettings, setProxySettings] = useState({}); const [activeTab, setActiveTab] = useState("models"); - const tableRef = useRef>(null); - const agentTableRef = useRef>(null); - const mcpTableRef = useRef>(null); useEffect(() => { const initializeAndFetch = async () => { @@ -1121,7 +1117,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded columns={publicModelHubColumns()} data={filteredData} isLoading={loading} - table={tableRef} defaultSorting={[{ id: "model_group", desc: false }]} /> @@ -1184,7 +1179,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded columns={publicAgentHubColumns()} data={filteredAgentData} isLoading={agentLoading} - table={agentTableRef} defaultSorting={[{ id: "name", desc: false }]} /> @@ -1248,7 +1242,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded columns={publicMCPHubColumns()} data={filteredMcpData} isLoading={mcpLoading} - table={mcpTableRef} defaultSorting={[{ id: "server_name", desc: false }]} /> diff --git a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx index 6d43e0af057..9dbe04bffb1 100644 --- a/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/templates/model_dashboard.tsx @@ -1,65 +1,74 @@ -import React, { useState, useEffect, useRef, useMemo } from "react"; import { Card, - Title, + Col, + Grid, Subtitle, Table, - TableHead, - TableRow, - TableHeaderCell, - TableCell, TableBody, + TableCell, + TableHead, + TableHeaderCell, + TableRow, Text, - Grid, - Col, + Title, } from "@tremor/react"; +import React, { useEffect, useMemo, useRef, useState } from "react"; import { CredentialItem, credentialListCall, CredentialsResponse } from "../networking"; import { handleAddModelSubmit } from "../add_model/handle_add_model_submit"; import CredentialsPanel from "@/components/model_add/credentials"; -import { getDisplayModelName } from "../view_model/model_name_display"; -import { TabPanel, TabPanels, TabGroup, TabList, Tab, Icon } from "@tremor/react"; -import { Select, SelectItem, DateRangePickerValue } from "@tremor/react"; -import UsageDatePicker from "../shared/usage_date_picker"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { FilterIcon, RefreshIcon } from "@heroicons/react/outline"; +import { + AreaChart, + BarChart, + Button, + DateRangePickerValue, + Icon, + Select, + SelectItem, + Tab, + TabGroup, + TabList, + TabPanel, + TabPanels, +} from "@tremor/react"; +import type { UploadProps } from "antd"; +import { Form, InputNumber, Popover, Typography } from "antd"; +import AddModelTab from "../add_model/add_model_tab"; +import { Team } from "../key_team_helpers/key_list"; +import ModelInfoView from "../model_info_view"; +import TimeToFirstToken from "../model_metrics/time_to_first_token"; import { - modelInfoCall, - modelCostMap, - healthCheckCall, - modelMetricsCall, - streamingModelMetricsCall, - modelExceptionsCall, - modelMetricsSlowResponsesCall, - getCallbacksCall, - setCallbacksCall, - modelSettingsCall, adminGlobalActivityExceptions, adminGlobalActivityExceptionsPerDeployment, allEndUsersCall, + getCallbacksCall, + healthCheckCall, + modelCostMap, + modelExceptionsCall, + modelInfoCall, + modelMetricsCall, + modelMetricsSlowResponsesCall, + modelSettingsCall, + setCallbacksCall, + streamingModelMetricsCall, } from "../networking"; -import { BarChart, AreaChart } from "@tremor/react"; -import { Popover, Form, InputNumber } from "antd"; -import { Button } from "@tremor/react"; -import { Typography } from "antd"; -import { RefreshIcon, FilterIcon } from "@heroicons/react/outline"; -import { InfoCircleOutlined } from "@ant-design/icons"; -import type { UploadProps } from "antd"; -import TimeToFirstToken from "../model_metrics/time_to_first_token"; -import { Team } from "../key_team_helpers/key_list"; +import { getPlaceholder, getProviderModels, provider_map, Providers } from "../provider_info_helpers"; +import UsageDatePicker from "../shared/usage_date_picker"; import TeamInfoView from "../team/team_info"; -import { Providers, provider_map, getPlaceholder, getProviderModels } from "../provider_info_helpers"; -import ModelInfoView from "../model_info_view"; -import AddModelTab from "../add_model/add_model_tab"; +import { getDisplayModelName } from "../view_model/model_name_display"; -import { ModelDataTable } from "../model_dashboard/table"; -import { columns } from "../molecules/models/columns"; -import PriceDataReload from "../price_data_reload"; -import HealthCheckComponent from "../model_dashboard/HealthCheckComponent"; -import PassThroughSettings from "../pass_through_settings"; -import ModelGroupAliasSettings from "../model_group_alias_settings"; import { all_admin_roles } from "@/utils/roles"; -import { Table as TableInstance, PaginationState } from "@tanstack/react-table"; +import { PaginationState } from "@tanstack/react-table"; +import HealthCheckComponent from "../model_dashboard/HealthCheckComponent"; +import { ModelDataTable } from "../model_dashboard/table"; +import ModelGroupAliasSettings from "../model_group_alias_settings"; +import { columns } from "../molecules/models/columns"; import NotificationsManager from "../molecules/notifications_manager"; +import PassThroughSettings from "../pass_through_settings"; +import PriceDataReload from "../price_data_reload"; interface ModelDashboardProps { accessToken: string | null; @@ -196,7 +205,6 @@ const OldModelDashboard: React.FC = ({ const [isDropdownOpen, setIsDropdownOpen] = useState(false); const [expandedRows, setExpandedRows] = useState>(new Set()); const dropdownRef = useRef(null); - const tableRef = useRef>(null); // Pagination state const [pagination, setPagination] = useState({ @@ -1325,7 +1333,6 @@ const OldModelDashboard: React.FC = ({ )} data={paginatedData} isLoading={false} - table={tableRef} /> From 85827aa217962d415ff72ffee41e650dfedb3a27 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 24 Dec 2025 17:12:36 -0800 Subject: [PATCH 087/330] Resize columns working --- .../src/components/all_keys_table.tsx | 75 ++++++++++++++++++- 1 file changed, 71 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/all_keys_table.tsx index a915fe06179..210ce09fa34 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/ui/litellm-dashboard/src/components/all_keys_table.tsx @@ -1,6 +1,6 @@ "use client"; import React, { useEffect, useState } from "react"; -import { ColumnDef } from "@tanstack/react-table"; +import { ColumnDef, ColumnResizeMode, ColumnResizeDirection } from "@tanstack/react-table"; import { Select, SelectItem } from "@tremor/react"; import { Button } from "@tremor/react"; import KeyInfoView from "./templates/key_info_view"; @@ -125,6 +125,8 @@ export function AllKeysTable({ }: AllKeysTableProps) { const [selectedKeyId, setSelectedKeyId] = useState(null); const [userList, setUserList] = useState([]); + const [columnResizeMode, setColumnResizeMode] = React.useState("onChange"); + const [columnResizeDirection, setColumnResizeDirection] = React.useState("ltr"); const [sorting, setSorting] = React.useState(() => { if (currentSort) { return [ @@ -184,6 +186,7 @@ export function AllKeysTable({ { id: "expander", header: () => null, + size: 40, cell: ({ row }) => row.getCanExpand() ? ( + )} + + + + {isSSOConfigured ? ( + renderSSOSettings() + ) : ( + setIsAddModalVisible(true)} /> + )} + + + setIsDeleteModalVisible(false)} + onSuccess={() => refetch()} + accessToken={accessToken} + /> + + setIsAddModalVisible(false)} + onSuccess={() => { + setIsAddModalVisible(false); + refetch(); + }} + accessToken={accessToken} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx new file mode 100644 index 00000000000..fc315493a54 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx @@ -0,0 +1,30 @@ +import { Empty, Typography, Button } from "antd"; + +const { Title, Paragraph } = Typography; + +interface SSOSettingsEmptyPlaceholderProps { + onAdd: () => void; +} + +export default function SSOSettingsEmptyPlaceholder({ onAdd }: SSOSettingsEmptyPlaceholderProps) { + return ( +
+ + No SSO Configuration Found + + Configure Single Sign-On (SSO) to enable seamless authentication for your team members using your identity + provider. + +
+ } + > + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx index 4ddd5cd5d1f..6af5a226da6 100644 --- a/ui/litellm-dashboard/src/components/admins.tsx +++ b/ui/litellm-dashboard/src/components/admins.tsx @@ -55,6 +55,7 @@ import { getSSOSettings, } from "./networking"; import UISettings from "./Settings/AdminSettings/UISettings/UISettings"; +import SSOSettings from "./Settings/AdminSettings/SSOSettings/SSOSettings"; const AdminPanel: React.FC = ({ searchParams, @@ -496,11 +497,15 @@ const AdminPanel: React.FC = ({ Go to 'Internal Users' page to add other admins. + SSO Settings Security Settings SCIM UI Settings + + + ✨ Security Settings From 546fba98498d8019d8ba911e2a4ee62dd35df1a7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 18:28:10 -0800 Subject: [PATCH 224/330] tests --- .../Modals/AddSSOSettingsModal.test.tsx | 26 +++++++++++++ .../Modals/DeleteSSOSettingsModal.test.tsx | 20 ++++++++++ .../SSOSettings/SSOSettings.test.tsx | 37 +++++++++++++++++++ .../SSOSettingsEmptyPlaceholder.test.tsx | 14 +++++++ 4 files changed, 97 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..13363a11643 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -0,0 +1,26 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import AddSSOSettingsModal from "./AddSSOSettingsModal"; + +// Mock networking functions +vi.mock("@/components/networking", () => ({ + updateSSOSettings: vi.fn(), +})); + +// Mock error utils +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn((error) => error?.message || "Unknown error"), +})); + +describe("AddSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + render(); + + expect(screen.getByText("SSO Provider")).toBeInTheDocument(); + expect(screen.getByText("Cancel")).toBeInTheDocument(); + expect(screen.getAllByText("Add SSO")).toHaveLength(2); // Title and button + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..ef6ec6c7055 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx @@ -0,0 +1,20 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import DeleteSSOSettingsModal from "./DeleteSSOSettingsModal"; + +describe("DeleteSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + render( + , + ); + + expect(screen.getByText("Confirm Clear SSO Settings")).toBeInTheDocument(); + expect( + screen.getByText("Are you sure you want to clear all SSO settings? This action cannot be undone."), + ).toBeInTheDocument(); + expect(screen.getByText("Users will no longer be able to login using SSO after this change.")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx new file mode 100644 index 00000000000..5e7908a872b --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx @@ -0,0 +1,37 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettings from "./SSOSettings"; + +// Mock the useSSOSettings hook +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: () => ({ + data: null, + refetch: vi.fn(), + }), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +describe("SSOSettings", () => { + it("should render", () => { + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(screen.getByText("SSO Configuration")).toBeInTheDocument(); + expect(screen.getByText("Manage Single Sign-On authentication settings")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx new file mode 100644 index 00000000000..6676ba1c2c9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx @@ -0,0 +1,14 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; + +describe("SSOSettingsEmptyPlaceholder", () => { + it("should render", () => { + const onAdd = vi.fn(); + + render(); + + expect(screen.getByText("No SSO Configuration Found")).toBeInTheDocument(); + expect(screen.getByText("Configure SSO")).toBeInTheDocument(); + }); +}); From f4c712506dc9b139b2e0b0e46865d3fd3a109121 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 18:49:01 -0800 Subject: [PATCH 225/330] Unit tests to increase test coverage --- .../hooks/uiSettings/useUISettings.test.ts | 185 +++++++ .../playground/chat_ui/EndpointUtils.test.tsx | 219 ++++++++ .../prompt_editor_view/ToolsCard.test.tsx | 82 +++ .../VersionHistorySidePanel.test.tsx | 473 ++++++++++++++++++ .../prompts/prompt_editor_view/utils.test.ts | 444 ++++++++++++++++ 5 files changed, 1403 insertions(+) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts create mode 100644 ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx create mode 100644 ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts new file mode 100644 index 00000000000..785f003d2f8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts @@ -0,0 +1,185 @@ +import { getUiSettings } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useUISettings } from "./useUISettings"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getUiSettings: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("../useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockUISettings: Record = { + theme: "dark", + language: "en", + notifications: true, + dashboard_layout: "compact", + api_keys_visible: false, +}; + +describe("useUISettings", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return UI settings data when query is successful", async () => { + // Mock successful API call + (getUiSettings as any).mockResolvedValue(mockUISettings); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUISettings); + expect(result.current.error).toBeNull(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getUiSettings fails", async () => { + const errorMessage = "Failed to fetch UI settings"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getUiSettings as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is empty string", async () => { + // Mock empty accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: "", + userRole: "Admin", + userId: "test-user-id", + token: "", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should return empty object when API returns empty settings", async () => { + // Mock API returning empty object + (getUiSettings as any).mockResolvedValue({}); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({}); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getUiSettings as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx new file mode 100644 index 00000000000..6eb481381ac --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx @@ -0,0 +1,219 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ModelGroup } from "../llm_calls/fetch_models"; +import { determineEndpointType } from "./EndpointUtils"; +import { EndpointType } from "./mode_endpoint_mapping"; + +// Mock the getEndpointType function +vi.mock("./mode_endpoint_mapping", () => ({ + EndpointType: { + IMAGE: "image", + VIDEO: "video", + CHAT: "chat", + RESPONSES: "responses", + IMAGE_EDITS: "image_edits", + ANTHROPIC_MESSAGES: "anthropic_messages", + EMBEDDINGS: "embeddings", + SPEECH: "speech", + TRANSCRIPTION: "transcription", + A2A_AGENTS: "a2a_agents", + }, + getEndpointType: vi.fn(), + ModelMode: { + AUDIO_SPEECH: "audio_speech", + AUDIO_TRANSCRIPTION: "audio_transcription", + IMAGE_GENERATION: "image_generation", + VIDEO_GENERATION: "video_generation", + CHAT: "chat", + RESPONSES: "responses", + IMAGE_EDITS: "image_edits", + ANTHROPIC_MESSAGES: "anthropic_messages", + EMBEDDING: "embedding", + }, +})); + +// Import the mocked function +import { getEndpointType } from "./mode_endpoint_mapping"; + +describe("determineEndpointType", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should return the correct endpoint type when model is found and has a valid mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + { + model_group: "dall-e-3", + mode: "image_generation", + }, + ]; + + // Mock getEndpointType to return IMAGE for image_generation mode + vi.mocked(getEndpointType).mockReturnValue(EndpointType.IMAGE); + + const result = determineEndpointType("dall-e-3", mockModelInfo); + + expect(getEndpointType).toHaveBeenCalledWith("image_generation"); + expect(result).toBe(EndpointType.IMAGE); + }); + + it("should return CHAT endpoint type when model is found but has no mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + // No mode property + }, + ]; + + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should return CHAT endpoint type when model is not found in modelInfo", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + ]; + + const result = determineEndpointType("non-existent-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should return CHAT endpoint type when modelInfo array is empty", () => { + const mockModelInfo: ModelGroup[] = []; + + const result = determineEndpointType("any-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle different mode types correctly", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "tts-model", + mode: "audio_speech", + }, + { + model_group: "whisper-model", + mode: "audio_transcription", + }, + { + model_group: "embedding-model", + mode: "embedding", + }, + { + model_group: "video-model", + mode: "video_generation", + }, + ]; + + // Test speech mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.SPEECH); + const speechResult = determineEndpointType("tts-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("audio_speech"); + expect(speechResult).toBe(EndpointType.SPEECH); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test transcription mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.TRANSCRIPTION); + const transcriptionResult = determineEndpointType("whisper-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("audio_transcription"); + expect(transcriptionResult).toBe(EndpointType.TRANSCRIPTION); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test embedding mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.EMBEDDINGS); + const embeddingResult = determineEndpointType("embedding-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("embedding"); + expect(embeddingResult).toBe(EndpointType.EMBEDDINGS); + + // Reset mock for next test + vi.clearAllMocks(); + + // Test video mode + vi.mocked(getEndpointType).mockReturnValueOnce(EndpointType.VIDEO); + const videoResult = determineEndpointType("video-model", mockModelInfo); + expect(getEndpointType).toHaveBeenCalledWith("video_generation"); + expect(videoResult).toBe(EndpointType.VIDEO); + }); + + it("should prioritize the first matching model when there are duplicates", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "gpt-3.5-turbo", + mode: "chat", + }, + { + model_group: "gpt-3.5-turbo", + mode: "image_generation", // Different mode for same model name + }, + ]; + + vi.mocked(getEndpointType).mockReturnValue(EndpointType.CHAT); + + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).toHaveBeenCalledWith("chat"); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle models with undefined mode property explicitly set", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "test-model", + mode: undefined, + }, + ]; + + const result = determineEndpointType("test-model", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle models with empty string mode", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "test-model", + mode: "", + }, + ]; + + const result = determineEndpointType("test-model", mockModelInfo); + + // Empty string is falsy, so getEndpointType should not be called + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); + + it("should handle case-sensitive model group matching", () => { + const mockModelInfo: ModelGroup[] = [ + { + model_group: "GPT-3.5-TURBO", + mode: "chat", + }, + ]; + + vi.mocked(getEndpointType).mockReturnValue(EndpointType.CHAT); + + // Test with different case - should not match + const result = determineEndpointType("gpt-3.5-turbo", mockModelInfo); + + expect(getEndpointType).not.toHaveBeenCalled(); + expect(result).toBe(EndpointType.CHAT); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx new file mode 100644 index 00000000000..742b8aa37fc --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx @@ -0,0 +1,82 @@ +import { act, fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import ToolsCard from "./ToolsCard"; +import { Tool } from "./types"; + +describe("ToolsCard", () => { + const mockTools: Tool[] = [ + { + name: "Calculator", + description: "Performs mathematical calculations", + json: '{"type": "function", "function": {"name": "calculate"}}', + }, + { + name: "Weather API", + description: "Gets current weather information", + json: '{"type": "function", "function": {"name": "get_weather"}}', + }, + ]; + + const defaultProps = { + tools: [] as Tool[], + onAddTool: vi.fn(), + onEditTool: vi.fn(), + onRemoveTool: vi.fn(), + }; + + it("should render the component", () => { + render(); + expect(screen.getByText("Tools")).toBeInTheDocument(); + }); + + it("should display no tools message when tools array is empty", () => { + render(); + expect(screen.getByText("No tools added")).toBeInTheDocument(); + }); + + it("should render tools when provided", () => { + render(); + + expect(screen.getByText("Calculator")).toBeInTheDocument(); + expect(screen.getByText("Performs mathematical calculations")).toBeInTheDocument(); + expect(screen.getByText("Weather API")).toBeInTheDocument(); + expect(screen.getByText("Gets current weather information")).toBeInTheDocument(); + }); + + it("should call onAddTool when Add button is clicked", () => { + const mockOnAddTool = vi.fn(); + render(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /add/i })); + }); + + expect(mockOnAddTool).toHaveBeenCalledTimes(1); + }); + + it("should call onEditTool with correct index when Edit button is clicked", () => { + const mockOnEditTool = vi.fn(); + render(); + + const editButtons = screen.getAllByText("Edit"); + act(() => { + fireEvent.click(editButtons[0]); + }); + + expect(mockOnEditTool).toHaveBeenCalledWith(0); + expect(mockOnEditTool).toHaveBeenCalledTimes(1); + }); + + it("should call onRemoveTool with correct index when remove button is clicked", () => { + const mockOnRemoveTool = vi.fn(); + render(); + + const removeButtons = screen.getAllByRole("button", { name: "" }); + act(() => { + fireEvent.click(removeButtons[0]); + }); + + expect(mockOnRemoveTool).toHaveBeenCalledWith(0); + expect(mockOnRemoveTool).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx new file mode 100644 index 00000000000..b0346e03a2d --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx @@ -0,0 +1,473 @@ +import { describe, it, expect, vi, beforeEach, afterEach, type Mock } from "vitest"; +import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; +import VersionHistorySidePanel from "./VersionHistorySidePanel"; +import { getPromptVersions } from "../../networking"; +import type { PromptSpec } from "../../networking"; + +// Mock the networking function +vi.mock("../../networking", () => ({ + getPromptVersions: vi.fn(), +})); + +const mockGetPromptVersions = getPromptVersions as Mock; + +// Mock Ant Design components that might need special handling +vi.mock("antd", async () => { + const actual = await vi.importActual("antd"); + return { + ...actual, + Drawer: ({ children, title, onClose, open, width, placement, mask, maskClosable }: any) => ( +
+
{title}
+ +
{children}
+
+ ), + List: ({ children, dataSource, renderItem }: any) => ( +
{dataSource?.map((item: any, index: number) => renderItem(item, index))}
+ ), + Skeleton: ({ active }: any) => ( +
+ Loading... +
+ ), + Tag: ({ children, color, className }: any) => ( + + {children} + + ), + Typography: { + Text: ({ children, type, className }: any) => ( + + {children} + + ), + }, + }; +}); + +describe("VersionHistorySidePanel", () => { + // Mock data + const mockPromptVersions: PromptSpec[] = [ + { + prompt_id: "test-prompt.v2", + litellm_params: { prompt_id: "test-prompt.v2" }, + prompt_info: { prompt_type: "db" }, + version: 2, + created_at: "2024-01-15T10:30:00Z", + }, + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + version: 1, + created_at: "2024-01-10T09:00:00Z", + }, + { + prompt_id: "test-prompt.v3", + litellm_params: { prompt_id: "test-prompt.v3" }, + prompt_info: { prompt_type: "config" }, + version: 3, + created_at: "2024-01-20T14:15:00Z", + }, + ]; + + const mockPromptVersionsWithoutExplicitVersion = [ + { + prompt_id: "test-prompt.v2", + litellm_params: { prompt_id: "test-prompt.v2" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-10T09:00:00Z", + }, + ]; + + const defaultProps = { + isOpen: true, + onClose: vi.fn(), + accessToken: "test-token", + promptId: "test-prompt.v2", + activeVersionId: "test-prompt.v2", + onSelectVersion: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + // Mock successful response by default + mockGetPromptVersions.mockResolvedValue({ + prompts: mockPromptVersions, + }); + }); + + afterEach(() => { + vi.clearAllTimers(); + }); + + describe("Component Rendering", () => { + it("should render the component with drawer", async () => { + await act(async () => { + render(); + }); + expect(screen.getByTestId("drawer")).toBeInTheDocument(); + expect(screen.getByText("Version History")).toBeInTheDocument(); + }); + + it("should not render when isOpen is false", async () => { + await act(async () => { + render(); + }); + // The drawer should still be rendered but with open=false + const drawer = screen.getByTestId("drawer"); + expect(drawer).toHaveAttribute("data-open", "false"); + }); + + it("should show loading skeleton initially", async () => { + // Mock a delayed response to show loading state + mockGetPromptVersions.mockImplementationOnce( + () => new Promise((resolve) => setTimeout(() => resolve({ prompts: mockPromptVersions }), 100)), + ); + + render(); + expect(screen.getByTestId("skeleton")).toBeInTheDocument(); + + // Wait for loading to complete + await waitFor(() => { + expect(screen.queryByTestId("skeleton")).not.toBeInTheDocument(); + }); + }); + + it("should show empty state when no versions are available", async () => { + mockGetPromptVersions.mockResolvedValueOnce({ prompts: [] }); + + render(); + + await waitFor(() => { + expect(screen.getByText("No version history available.")).toBeInTheDocument(); + }); + }); + + it("should render version list when data is loaded", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v1")).toBeInTheDocument(); + expect(screen.getByText("v3")).toBeInTheDocument(); + }); + + // Check that Latest tag is shown for the first item + const latestTags = screen.getAllByText("Latest"); + expect(latestTags.length).toBeGreaterThan(0); + + // Check Active tag is shown for the active version + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + describe("Version Selection and Highlighting", () => { + it("should highlight the active version correctly", async () => { + render(); + + await waitFor(() => { + const versionItems = screen.getAllByTestId("tag"); + // Should have Active tag for the selected version + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should highlight the latest version when no activeVersionId is provided", async () => { + render(); + + await waitFor(() => { + const latestTags = screen.getAllByText("Latest"); + expect(latestTags.length).toBeGreaterThan(0); + }); + }); + + it("should call onSelectVersion when a version is clicked", async () => { + const mockOnSelectVersion = vi.fn(); + render(); + + await waitFor(() => { + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + + const versionItem = screen.getByText("v1").closest("div"); + expect(versionItem).toBeInTheDocument(); + + act(() => { + fireEvent.click(versionItem!); + }); + + expect(mockOnSelectVersion).toHaveBeenCalledWith(mockPromptVersions[1]); + }); + }); + + describe("Version Number Extraction", () => { + it("should extract version from explicit version field", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v3")).toBeInTheDocument(); + }); + }); + + it("should extract version from prompt_id with .v suffix", async () => { + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: mockPromptVersionsWithoutExplicitVersion, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + }); + + it("should extract version from prompt_id with _v suffix", async () => { + const versionsWithUnderscore = [ + { + prompt_id: "test-prompt_v2", + litellm_params: { prompt_id: "test-prompt_v2" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithUnderscore, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v2")).toBeInTheDocument(); + }); + }); + + it("should default to v1 when no version info is available", async () => { + const versionWithoutVersionInfo = [ + { + prompt_id: "test-prompt", + litellm_params: { prompt_id: "test-prompt" }, + prompt_info: { prompt_type: "db" }, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionWithoutVersionInfo, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("v1")).toBeInTheDocument(); + }); + }); + }); + + describe("Date Formatting", () => { + it("should format dates correctly", async () => { + render(); + + await waitFor(() => { + // Check that dates are displayed (format: YYYY-MM-DD HH:MM:SS) + const dateElements = screen.getAllByTestId("text"); + const dateText = dateElements.find((el) => el.textContent?.match(/\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}/)); + expect(dateText).toBeTruthy(); + }); + }); + + it("should show dash for missing dates", async () => { + const versionsWithoutDates = [ + { + prompt_id: "test-prompt.v1", + litellm_params: { prompt_id: "test-prompt.v1" }, + prompt_info: { prompt_type: "db" }, + version: 1, + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithoutDates, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("-")).toBeInTheDocument(); + }); + }); + }); + + describe("Prompt Type Display", () => { + it("should show 'Saved to Database' for db prompts", async () => { + render(); + + await waitFor(() => { + const dbTexts = screen.getAllByText("Saved to Database"); + expect(dbTexts.length).toBeGreaterThan(0); + }); + }); + + it("should show 'Config Prompt' for config prompts", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Config Prompt")).toBeInTheDocument(); + }); + }); + }); + + describe("Network Calls and Data Fetching", () => { + it("should call getPromptVersions with correct parameters", async () => { + render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "test-prompt"); + }); + }); + + it("should strip .v suffix from promptId when fetching versions", async () => { + render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "test-prompt"); + }); + }); + + it("should not fetch versions when isOpen is false", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should not fetch versions when accessToken is null", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should not fetch versions when promptId is not provided", () => { + render(); + + expect(getPromptVersions).not.toHaveBeenCalled(); + }); + + it("should refetch versions when props change", async () => { + const { rerender } = render(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledTimes(1); + }); + + rerender(); + + await waitFor(() => { + expect(getPromptVersions).toHaveBeenCalledTimes(2); + expect(getPromptVersions).toHaveBeenCalledWith("test-token", "different-prompt"); + }); + }); + }); + + describe("Error Handling", () => { + it("should handle network errors gracefully", async () => { + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockGetPromptVersions.mockRejectedValueOnce(new Error("Network error")); + + render(); + + await waitFor(() => { + expect(consoleSpy).toHaveBeenCalledWith("Error fetching prompt versions:", expect.any(Error)); + }); + + // Should show empty state when there's an error + expect(screen.getByText("No version history available.")).toBeInTheDocument(); + + consoleSpy.mockRestore(); + }); + }); + + describe("User Interactions", () => { + it("should call onClose when close button is clicked", () => { + const mockOnClose = vi.fn(); + render(); + + const drawer = screen.getByTestId("drawer"); + act(() => { + fireEvent.click(drawer); // Simulate close action + }); + + // Note: This test assumes the drawer handles close events. + // In a real scenario, you'd test the actual close trigger. + }); + + it("should prevent interaction with main content when drawer is open", () => { + render(); + + const drawer = screen.getByTestId("drawer"); + // The mask and maskClosable props are passed as boolean false to disable them + expect(drawer).toHaveAttribute("data-mask", "false"); + expect(drawer).toHaveAttribute("data-maskclosable", "false"); + }); + }); + + describe("Edge Cases", () => { + it("should handle activeVersionId with .v suffix correctly", async () => { + render(); + + await waitFor(() => { + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should handle activeVersionId with _v suffix correctly", async () => { + const versionsWithUnderscore = [ + { + prompt_id: "test-prompt_v2", + litellm_params: { prompt_id: "test-prompt_v2" }, + prompt_info: { prompt_type: "db" }, + version: 2, + created_at: "2024-01-15T10:30:00Z", + }, + ]; + + mockGetPromptVersions.mockResolvedValueOnce({ + prompts: versionsWithUnderscore, + }); + + render(); + + await waitFor(() => { + expect(screen.getByText("Active")).toBeInTheDocument(); + }); + }); + + it("should sort versions correctly with version field", async () => { + // The component doesn't explicitly sort, but we can verify the order from the API response + render(); + + await waitFor(() => { + const versionElements = screen.getAllByTestId("tag"); + // Verify versions are displayed as they come from the API + expect(screen.getByText("v2")).toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts new file mode 100644 index 00000000000..59fcaccb639 --- /dev/null +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts @@ -0,0 +1,444 @@ +import { describe, expect, it } from "vitest"; +import { PromptType } from "./types"; +import { + convertToDotPrompt, + extractVariables, + getVersionNumber, + parseExistingPrompt, + stripVersionFromPromptId, +} from "./utils"; + +describe("extractVariables", () => { + it("should extract variables from messages", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello {{name}}, how are you?" }, + { role: "assistant", content: "I am fine {{name}}" }, + ], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["name"]); + }); + + it("should extract variables from developer message", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are {{role}} assistant", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["role"]); + }); + + it("should extract variables from both messages and developer message", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are {{role}} assistant", + messages: [ + { role: "user", content: "Hello {{name}}" }, + { role: "assistant", content: "Hi {{name}}, I am {{role}}" }, + ], + }; + + const result = extractVariables(prompt); + expect(result.sort()).toEqual(["name", "role"].sort()); + }); + + it("should return empty array when no variables present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are an assistant", + messages: [{ role: "user", content: "Hello world" }], + }; + + const result = extractVariables(prompt); + expect(result).toEqual([]); + }); + + it("should handle duplicate variables", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello {{name}}" }, + { role: "assistant", content: "Hi {{name}} again" }, + ], + }; + + const result = extractVariables(prompt); + expect(result).toEqual(["name"]); + }); +}); + +describe("convertToDotPrompt", () => { + it("should convert basic prompt to dot prompt format", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello world" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("---"); + expect(result).toContain("model: gpt-4"); + expect(result).toContain("input:"); + expect(result).toContain("schema:"); + expect(result).toContain("output:"); + expect(result).toContain("format: text"); + expect(result).toContain("User: Hello world"); + }); + + it("should include config parameters when set", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: { + temperature: 0.7, + max_tokens: 100, + top_p: 0.9, + }, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("temperature: 0.7"); + expect(result).toContain("max_tokens: 100"); + expect(result).toContain("top_p: 0.9"); + }); + + it("should include input schema with variables", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [{ role: "user", content: "Hello {{name}}" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("input:"); + expect(result).toContain("schema:"); + expect(result).toContain("name: string"); + }); + + it("should include developer message when present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "You are a helpful assistant", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("Developer: You are a helpful assistant"); + }); + + it("should include tools when present", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [ + { + name: "get_weather", + description: "Get weather information", + json: '{"type": "function", "function": {"name": "get_weather"}}', + }, + ], + developerMessage: "", + messages: [{ role: "user", content: "Hello" }], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("tools:"); + expect(result).toContain('{"type":"function","function":{"name":"get_weather"}}'); + }); + + it("should handle multiple messages with different roles", () => { + const prompt: PromptType = { + name: "test", + model: "gpt-4", + config: {}, + tools: [], + developerMessage: "", + messages: [ + { role: "user", content: "Hello" }, + { role: "assistant", content: "Hi there" }, + { role: "user", content: "How are you?" }, + ], + }; + + const result = convertToDotPrompt(prompt); + expect(result).toContain("User: Hello"); + expect(result).toContain("Assistant: Hi there"); + expect(result).toContain("User: How are you?"); + }); +}); + +describe("parseExistingPrompt", () => { + it("should parse basic dotprompt content", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello world`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.name).toBe("test-prompt"); + expect(result.model).toBe("gpt-4"); + expect(result.messages).toEqual([{ role: "user", content: "Hello world" }]); + }); + + it("should parse with config parameters", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +temperature: 0.7 +max_tokens: 100 +top_p: 0.9 +input: + schema: +output: + format: text +--- + +User: Hello`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.config.temperature).toBe(0.7); + expect(result.config.max_tokens).toBe(100); + expect(result.config.top_p).toBe(0.9); + }); + + it("should parse with developer message", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +Developer: You are a helpful assistant + +User: Hello`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.developerMessage).toBe("You are a helpful assistant"); + }); + + it("should parse multiple messages", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello +How are you? + +Assistant: I am fine +Thank you for asking + +User: Great!`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.messages).toEqual([ + { role: "user", content: "Hello\nHow are you?" }, + { role: "assistant", content: "I am fine\nThank you for asking" }, + { role: "user", content: "Great!" }, + ]); + }); + + it("should handle prompt with version suffix", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +User: Hello`, + }, + prompt_id: "test-prompt.v2", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.name).toBe("test-prompt"); + }); + + it("should throw error when no dotprompt_content", () => { + const apiResponse = { + prompt_spec: { + litellm_params: {}, + }, + }; + + expect(() => parseExistingPrompt(apiResponse)).toThrow("No dotprompt_content found in API response"); + }); + + it("should throw error for invalid dotprompt format", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: "invalid format", + }, + }, + }; + + expect(() => parseExistingPrompt(apiResponse)).toThrow("Invalid dotprompt format"); + }); + + it("should provide default values when parsing fails", () => { + const apiResponse = { + prompt_spec: { + litellm_params: { + dotprompt_content: `--- +model: gpt-4 +input: + schema: +output: + format: text +--- + +`, + }, + prompt_id: "test-prompt", + }, + }; + + const result = parseExistingPrompt(apiResponse); + expect(result.messages).toEqual([ + { role: "user", content: "Enter task specifics. Use {{template_variables}} for dynamic inputs" }, + ]); + }); +}); + +describe("getVersionNumber", () => { + it("should return '1' for undefined promptId", () => { + const result = getVersionNumber(undefined); + expect(result).toBe("1"); + }); + + it("should return '1' for promptId without version", () => { + const result = getVersionNumber("test-prompt"); + expect(result).toBe("1"); + }); + + it("should extract version with dot separator", () => { + const result = getVersionNumber("test-prompt.v2"); + expect(result).toBe("2"); + }); + + it("should extract version with underscore separator", () => { + const result = getVersionNumber("test-prompt_v3"); + expect(result).toBe("3"); + }); + + it("should extract version with hyphen separator", () => { + const result = getVersionNumber("test-prompt-v4"); + expect(result).toBe("4"); + }); + + it("should extract multi-digit version", () => { + const result = getVersionNumber("test-prompt.v123"); + expect(result).toBe("123"); + }); +}); + +describe("stripVersionFromPromptId", () => { + it("should return empty string for undefined promptId", () => { + const result = stripVersionFromPromptId(undefined); + expect(result).toBe(""); + }); + + it("should return promptId unchanged when no version present", () => { + const result = stripVersionFromPromptId("test-prompt"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with dot separator", () => { + const result = stripVersionFromPromptId("test-prompt.v2"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with underscore separator", () => { + const result = stripVersionFromPromptId("test-prompt_v3"); + expect(result).toBe("test-prompt"); + }); + + it("should strip version with hyphen separator", () => { + const result = stripVersionFromPromptId("test-prompt-v4"); + expect(result).toBe("test-prompt"); + }); + + it("should strip multi-digit version", () => { + const result = stripVersionFromPromptId("test-prompt.v123"); + expect(result).toBe("test-prompt"); + }); +}); From e6da33dc4ac9f3c39c2b58cb113d2e4df62ec8f0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 2 Jan 2026 19:22:56 -0800 Subject: [PATCH 226/330] Remove modal in useful links --- .../AIHub/UsefulLinksManagement.test.tsx | 158 +++++++++++++++++- .../AIHub/UsefulLinksManagement.tsx | 90 +++++----- 2 files changed, 189 insertions(+), 59 deletions(-) diff --git a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx index 7b0651ad2d6..0a859ca95f8 100644 --- a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.test.tsx @@ -1,10 +1,9 @@ +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { getProxyBaseUrl, getPublicModelHubInfo, updateUsefulLinksCall } from "@/components/networking"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { Modal } from "antd"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import UsefulLinksManagement from "./UsefulLinksManagement"; -import NotificationsManager from "@/components/molecules/notifications_manager"; -import { getPublicModelHubInfo, updateUsefulLinksCall, getProxyBaseUrl } from "@/components/networking"; vi.mock("@/components/networking", () => ({ getPublicModelHubInfo: vi.fn(), @@ -25,8 +24,6 @@ const mockedUpdateUsefulLinksCall = vi.mocked(updateUsefulLinksCall); const mockedGetProxyBaseUrl = vi.mocked(getProxyBaseUrl); const mockedNotifications = vi.mocked(NotificationsManager); -let modalSuccessSpy: any; - describe("UsefulLinksManagement", () => { beforeEach(() => { mockedGetPublicModelHubInfo.mockResolvedValue({ @@ -37,11 +34,9 @@ describe("UsefulLinksManagement", () => { }); mockedUpdateUsefulLinksCall.mockResolvedValue({}); mockedGetProxyBaseUrl.mockReturnValue("https://proxy.example.com"); - modalSuccessSpy = vi.spyOn(Modal, "success").mockImplementation(() => ({ destroy: vi.fn() }) as any); }); afterEach(() => { - modalSuccessSpy.mockRestore(); vi.clearAllMocks(); }); @@ -108,4 +103,153 @@ describe("UsefulLinksManagement", () => { expect(mockedNotifications.success).toHaveBeenCalledWith("Link order saved successfully"); }); + + it("should display the Model Hub link", async () => { + render(); + + expect(await screen.findByRole("link", { name: /public model hub/i })).toBeInTheDocument(); + }); + + it("should edit a link when edit button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Should show input fields in edit mode + expect(screen.getByDisplayValue("Test Link")).toBeInTheDocument(); + expect(screen.getByDisplayValue("https://test.example.com")).toBeInTheDocument(); + }); + + it("should update a link when save is clicked in edit mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click save + await user.click(screen.getByRole("button", { name: /save/i })); + + await waitFor(() => + expect(mockedUpdateUsefulLinksCall).toHaveBeenCalledWith("token", { + "Updated Link": { url: "https://test.example.com", index: 0 }, + }), + ); + + expect(mockedNotifications.success).toHaveBeenCalledWith("Link updated successfully"); + }); + + it("should cancel editing when cancel button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click cancel + await user.click(screen.getByRole("button", { name: /cancel/i })); + + // Should go back to normal view + expect(screen.getByText("Test Link")).toBeInTheDocument(); + expect(screen.queryByDisplayValue("Updated Link")).not.toBeInTheDocument(); + }); + + it("should not move down the last item in rearrange mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "First Link": "https://first.example.com", + "Second Link": "https://second.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("First Link")).toBeInTheDocument()); + + // Enter rearrange mode + await user.click(screen.getByRole("button", { name: /rearrange order/i })); + + // Try to move down the last item (should not do anything) + const secondLinkMoveDownButton = screen.getByTestId("move-down-1-Second Link"); + await user.click(secondLinkMoveDownButton); + + // Links should remain in same order + const linksAfter = screen.getAllByText(/First Link|Second Link/); + expect(linksAfter[0]).toHaveTextContent("First Link"); + expect(linksAfter[1]).toHaveTextContent("Second Link"); + }); + + it("should expand and collapse the component", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => expect(screen.getByText("Link Management")).toBeInTheDocument()); + + // Initially expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + + // Click to collapse + await user.click(screen.getByText("Link Management")); + + // Should be collapsed + expect(screen.queryByText("Manage Existing Links")).not.toBeInTheDocument(); + + // Click to expand again + await user.click(screen.getByText("Link Management")); + + // Should be expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx index 220113ded13..c73eaf52384 100644 --- a/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx @@ -1,11 +1,11 @@ -import React, { useState, useEffect } from "react"; -import { Modal } from "antd"; -import { PlusCircleIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { isAdminRole } from "@/utils/roles"; -import { getPublicModelHubInfo, updateUsefulLinksCall, getProxyBaseUrl } from "../networking"; -import { Card, Title, Text, Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; -import NotificationsManager from "@/components/molecules/notifications_manager"; import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { isAdminRole } from "@/utils/roles"; +import { ChevronDownIcon, ChevronRightIcon, ExternalLinkIcon, PlusCircleIcon } from "@heroicons/react/outline"; +import { Card, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text, Title } from "@tremor/react"; +import Link from "next/link"; +import React, { useEffect, useState } from "react"; +import { getProxyBaseUrl, getPublicModelHubInfo, updateUsefulLinksCall } from "../networking"; interface UsefulLinksManagementProps { accessToken: string | null; @@ -102,32 +102,6 @@ const UsefulLinksManagement: React.FC = ({ accessTok }); await updateUsefulLinksCall(accessToken, linksObject); - // show success modal with public model hub link - Modal.success({ - title: "Links Saved Successfully", - content: ( -
-

- Your useful links have been saved and are now visible on the public model hub. -

-
-

View your updated model hub:

- - Open Public Model Hub → - -
-
- ), - width: 500, - okText: "Close", - maskClosable: true, - keyboard: true, - }); return true; } catch (error) { @@ -319,29 +293,41 @@ const UsefulLinksManagement: React.FC = ({ accessTok
Manage Existing Links - {!isRearranging ? ( - - ) : ( -
+ Public Model Hub + + + {!isRearranging ? ( - -
- )} + ) : ( +
+ + +
+ )} +
From 0aae5153b6b59f9bbe6f479e70e4b97eaafa8365 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 3 Jan 2026 16:06:07 +0530 Subject: [PATCH 227/330] docs: Clarify Bedrock AgentCore documentation (#18603) Co-authored-by: Cursor Agent --- docs/my-website/docs/providers/bedrock_agentcore.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md index 43df7f82519..e3e352f7ab6 100644 --- a/docs/my-website/docs/providers/bedrock_agentcore.md +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -11,6 +11,12 @@ Call Bedrock AgentCore in the OpenAI Request/Response format. | Provider Route on LiteLLM | `bedrock/agentcore/{AGENT_RUNTIME_ARN}` | | Provider Doc | [AWS Bedrock AgentCore ↗](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) | +:::info + +This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details. + +::: + ## Quick Start ### Model Format to LiteLLM From 87fe62229f4b8b5dddefa7c22521eb5662928ca1 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 3 Jan 2026 21:51:19 +0530 Subject: [PATCH 228/330] feat: Add adopters page and data structure (#18605) Co-authored-by: Cursor Agent --- docs/my-website/src/data/adopters/README.md | 88 +++++++++++++++++++ .../src/data/adopters/adopters.json | 8 ++ docs/my-website/src/data/adopters/index.js | 23 +++++ .../img/adopters/placeholder-company.svg | 8 ++ 4 files changed, 127 insertions(+) create mode 100644 docs/my-website/src/data/adopters/README.md create mode 100644 docs/my-website/src/data/adopters/adopters.json create mode 100644 docs/my-website/src/data/adopters/index.js create mode 100644 docs/my-website/static/img/adopters/placeholder-company.svg diff --git a/docs/my-website/src/data/adopters/README.md b/docs/my-website/src/data/adopters/README.md new file mode 100644 index 00000000000..61a5215f802 --- /dev/null +++ b/docs/my-website/src/data/adopters/README.md @@ -0,0 +1,88 @@ +# LiteLLM Adopters + +This directory contains data for organizations that use LiteLLM in production. + +## Adding Your Organization + +We've made it super easy to add your organization! Just follow the steps below. + +### Quick Add (Recommended) + +**[Edit adopters.json on GitHub →](https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json)** + +This will open the GitHub editor in your browser where you can: + +1. Add your organization's entry to the JSON array +2. Commit your changes +3. GitHub will automatically create a pull request for you! + +No need to clone the repository or set up a development environment. + +### JSON Format + +Add your organization to the array in `adopters.json`: + +```json +{ + "name": "Your Organization Name", + "logoUrl": "https://yoursite.com/logo.svg", + "url": "https://yourcompany.com", + "description": "Brief description of how you use LiteLLM (shown on hover)" +} +``` + +### Fields + +- **`name`** (required): Your organization's display name +- **`logoUrl`** (required): URL to your logo - can be either: + - External URL: `https://yoursite.com/logo.svg` (easiest!) + - Local path: `/img/adopters/your-logo.svg` (requires uploading logo file) +- **`url`** (optional): Your organization's website (makes the logo clickable) +- **`description`** (optional): Brief description shown when users hover over your logo + +### Logo Options + +#### Option 1: External URL (Easiest) + +Simply provide a direct link to your logo hosted anywhere: + +```json +"logoUrl": "https://yourcompany.com/assets/logo.svg" +``` + +#### Option 2: Local Logo (Better Performance) + +If you prefer to host the logo locally: + +1. Add your logo to `docs/my-website/static/img/adopters/your-company.svg` +2. Reference it as: `"logoUrl": "/img/adopters/your-company.svg"` + +**Logo Specifications:** + +- **Format**: SVG preferred (PNG also acceptable) +- **Dimensions**: 240x160px or similar 3:2 ratio recommended +- **Background**: Transparent or white background works best + +### Example + +```json +{ + "name": "Acme Corporation", + "logoUrl": "https://acme.com/logo.svg", + "url": "https://acme.com", + "description": "Using LiteLLM to route requests across 50+ LLM providers" +} +``` + +### Display Order + +Adopters are displayed alphabetically by organization name, so your position will be determined automatically. + +### Need Help? + +If you have questions about adding your organization: + +- Ask in [GitHub Discussions](https://github.com/BerriAI/litellm/discussions) +- Join our [Discord community](https://discord.com/invite/wuPM9dRgDw) + +Thank you for supporting LiteLLM! 🚅 diff --git a/docs/my-website/src/data/adopters/adopters.json b/docs/my-website/src/data/adopters/adopters.json new file mode 100644 index 00000000000..52319c149e2 --- /dev/null +++ b/docs/my-website/src/data/adopters/adopters.json @@ -0,0 +1,8 @@ +[ + { + "name": "Your Logo Here", + "logoUrl": "/img/adopters/placeholder-company.svg", + "description": "Add your organization to show support for LiteLLM", + "url": "https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json" + } +] diff --git a/docs/my-website/src/data/adopters/index.js b/docs/my-website/src/data/adopters/index.js new file mode 100644 index 00000000000..b1a242dcc33 --- /dev/null +++ b/docs/my-website/src/data/adopters/index.js @@ -0,0 +1,23 @@ +import adoptersData from './adopters.json'; + +/** + * @typedef {Object} Adopter + * @property {string} name - The organization's display name + * @property {string} logoUrl - URL to the organization's logo + * @property {string} [url] - The organization's website URL + * @property {string} [description] - Brief description shown on hover + */ + +/** + * List of organizations using LiteLLM + * @type {Adopter[]} + */ +export const adopters = adoptersData; + +/** + * Adopters sorted alphabetically by name + * @type {Adopter[]} + */ +export const sortedAdopters = [...adopters].sort((a, b) => + a.name.localeCompare(b.name) +); diff --git a/docs/my-website/static/img/adopters/placeholder-company.svg b/docs/my-website/static/img/adopters/placeholder-company.svg new file mode 100644 index 00000000000..937dffc6eaf --- /dev/null +++ b/docs/my-website/static/img/adopters/placeholder-company.svg @@ -0,0 +1,8 @@ + + + + + + Add Your Logo + Click to contribute + From bdd05475bca872f944a7cf704d20cd3fdb72e0cc Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Sat, 3 Jan 2026 15:39:00 -0300 Subject: [PATCH 229/330] fix: correct cost calculation when reasoning_tokens present without text_tokens (#18607) Fixes #18599 When OpenAI models (gpt-5-nano, o1-*, o3-*) and other providers return reasoning_tokens in completion_tokens_details but don't provide text_tokens, LiteLLM was incorrectly calculating costs using only reasoning_tokens, ignoring the remaining completion tokens. Changes: - Modified generic_cost_per_token() in llm_cost_calc/utils.py to calculate text_tokens as: completion_tokens - reasoning_tokens - audio_tokens - image_tokens when text_tokens is not explicitly provided - Added comprehensive test case test_reasoning_tokens_without_text_tokens_gpt5_nano() to verify all completion_tokens are billed correctly Example: - completion_tokens: 977 - reasoning_tokens: 768 - Before: only 768 tokens billed (21% less) - After: all 977 tokens billed correctly Affected models: - OpenAI: gpt-5-nano, o1-*, o3-* - Perplexity: sonar-reasoning* - Any model returning reasoning_tokens without text_tokens --- .../litellm_core_utils/llm_cost_calc/utils.py | 20 ++++++-- .../llm_cost_calc/test_llm_cost_calc_utils.py | 51 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e36d6d68367..cbc0763382c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -604,12 +604,22 @@ def generic_cost_per_token( reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] - # Only assume all tokens are text if there's NO breakdown at all - # If image_tokens, audio_tokens, or reasoning_tokens exist, respect text_tokens=0 + # Handle text_tokens calculation: + # 1. If text_tokens is explicitly provided and > 0, use it + # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder + # 3. If no breakdown at all, assume all completion_tokens are text_tokens has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 - if text_tokens == 0 and not has_token_breakdown: - text_tokens = usage.completion_tokens - is_text_tokens_total = True + if text_tokens == 0: + if has_token_breakdown: + # Calculate text tokens as remainder when we have a breakdown + # This handles cases like OpenAI's reasoning models where text_tokens isn't provided + text_tokens = max( + 0, usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens + ) + else: + # No breakdown at all, all tokens are text tokens + text_tokens = usage.completion_tokens + is_text_tokens_total = True ## TEXT COST completion_cost = float(text_tokens) * completion_base_cost diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 65e3dbec8bd..5ba78d9eed1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -809,3 +809,54 @@ def test_bedrock_anthropic_prompt_caching(): assert completion_cost >= 0 assert round(prompt_cost, 3) == 0.111 assert round(completion_cost, 5) == 0.00820 + + +def test_reasoning_tokens_without_text_tokens_gpt5_nano(): + """ + Test fix for GitHub issue #18599: + https://github.com/BerriAI/litellm/issues/18599 + + When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide + text_tokens, LiteLLM should calculate text_tokens as: + text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens + + This ensures ALL completion tokens are billed, not just reasoning tokens. + """ + model = "gpt-5-nano" + custom_llm_provider = "openai" + + # Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided + # completion_tokens: 977 total + # reasoning_tokens: 768 + # text_tokens: should be calculated as 977 - 768 = 209 + usage = Usage( + prompt_tokens=17, + completion_tokens=977, + total_tokens=994, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=768, + audio_tokens=0, + # text_tokens NOT provided - this is the key part of the bug + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + # gpt-5-nano pricing: $0.05/1M input, $0.40/1M output + expected_prompt_cost = 17 * 0.05 / 1_000_000 + expected_completion_cost = 977 * 0.40 / 1_000_000 # ALL tokens, not just reasoning + + assert abs(prompt_cost - expected_prompt_cost) < 1e-10, \ + f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}" + + assert abs(completion_cost - expected_completion_cost) < 1e-10, \ + f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}" + + # Verify it's NOT using only reasoning_tokens (the bug) + wrong_cost = 768 * 0.40 / 1_000_000 # Only reasoning tokens + assert abs(completion_cost - wrong_cost) > 1e-6, \ + "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!" From 969790c4631f9efcbed0c55e0fb185ba98a0aa77 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sun, 4 Jan 2026 00:10:07 +0530 Subject: [PATCH 230/330] Iam roles anywhere docs (#18559) * Add documentation for IAM Roles Anywhere Co-authored-by: krrishdholakia * Refactor Bedrock provider docs for IAM Roles Anywhere Co-authored-by: krrishdholakia --------- Co-authored-by: Cursor Agent --- docs/my-website/docs/providers/bedrock.md | 47 +++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 122554fe8a4..f1eed4b4d52 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -2208,6 +2208,53 @@ response = completion( | `aws_role_name` | `RoleArn` | The Amazon Resource Name (ARN) of the role to assume | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | | `aws_session_name` | `RoleSessionName` | An identifier for the assumed role session | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | +### IAM Roles Anywhere (On-Premise / External Workloads) + +[IAM Roles Anywhere](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/introduction.html) extends IAM roles to workloads **outside of AWS** (on-premise servers, edge devices, other clouds). It uses the same STS mechanism as regular IAM roles but authenticates via X.509 certificates instead of AWS credentials. + +**Setup**: Configure the [AWS Signing Helper](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/credential-helper.html) as a credential process in `~/.aws/config`: + +```ini +[profile litellm-roles-anywhere] +credential_process = aws_signing_helper credential-process \ + --certificate /path/to/certificate.pem \ + --private-key /path/to/private-key.pem \ + --trust-anchor-arn arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/abc123 \ + --profile-arn arn:aws:rolesanywhere:us-east-1:123456789012:profile/def456 \ + --role-arn arn:aws:iam::123456789012:role/MyBedrockRole +``` + +**Usage**: Reference the profile in LiteLLM: + + + + +```python +from litellm import completion + +response = completion( + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + aws_profile_name="litellm-roles-anywhere", +) +``` + + + + +```yaml +model_list: + - model_name: bedrock-claude + litellm_params: + model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0 + aws_profile_name: "litellm-roles-anywhere" +``` + + + + +See the [IAM Roles Anywhere Getting Started Guide](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/getting-started.html) for trust anchor and profile setup. + Make the bedrock completion call From a3503e59c227f5f1dd15e9d8910f67a6b7dca2a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mehmet=20Can=20=C5=9Eakiro=C4=9Flu?= <53798389+cansakiroglu@users.noreply.github.com> Date: Sat, 3 Jan 2026 21:52:50 +0300 Subject: [PATCH 231/330] Litellm feat helm lifecycle support (#18517) * feat(helm): add lifecycle hook support for helm * add tests --- .../litellm-helm/templates/deployment.yaml | 4 ++++ .../litellm-helm/tests/deployment_tests.yaml | 24 ++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 0dab2ec40e0..19fa0479091 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -182,6 +182,10 @@ spec: {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} {{- with .Values.extraContainers }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index f9c83966696..182a2362392 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -136,4 +136,26 @@ tests: path: spec.template.spec.containers[0].volumeMounts content: name: litellm-config - mountPath: /etc/litellm/ \ No newline at end of file + mountPath: /etc/litellm/ + - it: should work with lifecycle hooks + template: deployment.yaml + set: + lifecycle: + preStop: + exec: + command: + - /bin/sh + - -c + - echo "Container stopping" + asserts: + - exists: + path: spec.template.spec.containers[0].lifecycle + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[0] + value: /bin/sh + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[1] + value: -c + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[2] + value: echo "Container stopping" \ No newline at end of file From 3a4ebf173f637dc9aa1fdc037b27a3757e246851 Mon Sep 17 00:00:00 2001 From: Deepak Walia <58362408+dee-walia20@users.noreply.github.com> Date: Sun, 4 Jan 2026 00:35:53 +0530 Subject: [PATCH 232/330] fix(sap): honor allowed_openai_params in transform_request (#18432) --- litellm/llms/sap/chat/transformation.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 01ceb72c0de..e13abca59f8 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -203,9 +203,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): headers: dict, ) -> dict: supported_params = self.get_supported_openai_params(model) + # Include extra params that passed validation (e.g., thinking_config for Gemini models via allowed_openai_params) + extra_params = [k for k in optional_params if k not in supported_params and k not in {"tools", "model_version"}] + supported_params = supported_params + extra_params model_params = { k: v for k, v in optional_params.items() if k in supported_params } + model_version = optional_params.pop("model_version", "latest") template = [] for message in messages: From dc62cdb3009bff03a2282de1a981f29421e0383c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=82=86=E3=82=8A?= Date: Sun, 4 Jan 2026 03:07:52 +0800 Subject: [PATCH 233/330] fix: handle empty error objects in response conversion (#18493) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Some OpenAI-compatible providers (e.g., Apertis) return empty error objects even on successful responses. The previous check only verified that error was not None, causing spurious APIErrors. Now the code checks if the error object contains meaningful data: - For dict errors: non-empty message OR non-null code - For string errors: non-empty string - Other truthy values are still treated as errors Fixes #18407 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: yurekami Co-authored-by: Claude Opus 4.5 --- .../convert_dict_to_response.py | 46 ++++-- .../test_convert_dict_to_chat_completion.py | 156 ++++++++++++++++++ 2 files changed, 188 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 59d2a8a8dd0..bbe28e3ec2c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -445,25 +445,43 @@ def convert_to_model_response_object( # noqa: PLR0915 hidden_params["additional_headers"] = additional_headers ### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary + # Some OpenAI-compatible providers (e.g., Apertis) return empty error objects + # even on success. Only raise if the error contains meaningful data. if ( response_object is not None and "error" in response_object and response_object["error"] is not None ): - error_args = {"status_code": 422, "message": "Error in response object"} - if isinstance(response_object["error"], dict): - if "code" in response_object["error"]: - error_args["status_code"] = response_object["error"]["code"] - if "message" in response_object["error"]: - if isinstance(response_object["error"]["message"], dict): - message_str = json.dumps(response_object["error"]["message"]) - else: - message_str = str(response_object["error"]["message"]) - error_args["message"] = message_str - raised_exception = Exception() - setattr(raised_exception, "status_code", error_args["status_code"]) - setattr(raised_exception, "message", error_args["message"]) - raise raised_exception + error_obj = response_object["error"] + has_meaningful_error = False + + if isinstance(error_obj, dict): + # Check if error dict has non-empty message or non-null code + error_message = error_obj.get("message", "") + error_code = error_obj.get("code") + has_meaningful_error = bool(error_message) or error_code is not None + elif isinstance(error_obj, str): + # String error is meaningful if non-empty + has_meaningful_error = bool(error_obj) + else: + # Any other truthy value is considered meaningful + has_meaningful_error = True + + if has_meaningful_error: + error_args = {"status_code": 422, "message": "Error in response object"} + if isinstance(error_obj, dict): + if "code" in error_obj: + error_args["status_code"] = error_obj["code"] + if "message" in error_obj: + if isinstance(error_obj["message"], dict): + message_str = json.dumps(error_obj["message"]) + else: + message_str = str(error_obj["message"]) + error_args["message"] = message_str + raised_exception = Exception() + setattr(raised_exception, "status_code", error_args["status_code"]) + setattr(raised_exception, "message", error_args["message"]) + raise raised_exception try: if response_type == "completion" and ( diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 7e269f21451..c151150f634 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -903,3 +903,159 @@ def test_convert_to_model_response_object_with_thinking_content(): resp: ModelResponse = convert_to_model_response_object(**args) assert resp is not None assert resp.choices[0].message.reasoning_content is not None + + +def test_convert_to_model_response_object_with_empty_error_object(): + """ + Test that convert_to_model_response_object handles empty error objects gracefully. + + This is a regression test for issue #18407 where providers like Apertis return + empty error objects even on successful responses, causing spurious APIErrors. + + The error object structure: + { + "error": { + "message": "", + "type": "", + "param": "", + "code": null + } + } + """ + response_object = { + "model": "minimax-m2.1", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hey! I'm doing well, thanks for asking!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 49, + "completion_tokens": 87, + "total_tokens": 136, + }, + "error": { + "message": "", + "type": "", + "param": "", + "code": None, + }, + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.model == "minimax-m2.1" + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Hey! I'm doing well, thanks for asking!" + + +def test_convert_to_model_response_object_with_real_error(): + """ + Test that convert_to_model_response_object still raises for real errors. + + Ensures the empty error fix doesn't break legitimate error handling. + """ + response_object = { + "error": { + "message": "Rate limit exceeded", + "type": "rate_limit_error", + "param": None, + "code": 429, + }, + } + + with pytest.raises(Exception) as exc_info: + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + # The exception should have the error message + assert hasattr(exc_info.value, "message") + assert "Rate limit exceeded" in str(exc_info.value.message) + + +def test_convert_to_model_response_object_with_empty_dict_error(): + """ + Test that convert_to_model_response_object handles completely empty error dict. + """ + response_object = { + "model": "test-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + "error": {}, # Completely empty error object + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello!" + + +def test_convert_to_model_response_object_with_error_code_only(): + """ + Test that errors with only a code (no message) are still treated as real errors. + """ + response_object = { + "error": { + "message": "", + "code": 500, + }, + } + + with pytest.raises(Exception): + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) From 9ba27d85cee19e2cced4fbac4a2b68e7d7ae7dcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=82=86=E3=82=8A?= Date: Sun, 4 Jan 2026 03:08:32 +0800 Subject: [PATCH 234/330] feat(types): add output_text property to ResponsesAPIResponse (#18491) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the output_text convenience property to ResponsesAPIResponse that aggregates all output_text items from the output list, matching the OpenAI SDK's Response.output_text behavior. The property iterates through output items, collects text content from message-type outputs, and returns them concatenated into a single string. Returns empty string if no output_text content exists. Handles both dict and Pydantic model access patterns for compatibility with different output formats. Fixes #18470 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: yurekami Co-authored-by: Claude Opus 4.5 --- litellm/types/llms/openai.py | 33 +++++ .../types/llms/test_types_llms_openai.py | 134 ++++++++++++++++++ 2 files changed, 167 insertions(+) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index ceeae958a80..c2912558cab 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1197,6 +1197,39 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) + @property + def output_text(self) -> str: + """ + Convenience property that aggregates all `output_text` items from the `output` list. + + If no `output_text` content blocks exist, then an empty string is returned. + + This matches the OpenAI SDK's Response.output_text behavior. + """ + texts: List[str] = [] + for output_item in self.output: + # Handle both dict and object access patterns + if isinstance(output_item, dict): + item_type = output_item.get("type") + content = output_item.get("content", []) + else: + item_type = getattr(output_item, "type", None) + content = getattr(output_item, "content", []) + + if item_type == "message": + for content_item in content: + if isinstance(content_item, dict): + content_type = content_item.get("type") + text = content_item.get("text", "") + else: + content_type = getattr(content_item, "type", None) + text = getattr(content_item, "text", "") or "" + + if content_type == "output_text": + texts.append(text) + + return "".join(texts) + class ResponsesAPIStreamEvents(str, Enum): """ diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 05dec06d469..87cc9586665 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -35,3 +35,137 @@ def test_output_item_added_event(): assert event.sequence_number == 4 assert event.output_index == 1 assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" From 37c908caf9777a757f813596ec04897aa995170f Mon Sep 17 00:00:00 2001 From: Lu Date: Sun, 4 Jan 2026 03:13:22 +0800 Subject: [PATCH 235/330] google genai adapter inline data support (#18477) * support inline data * add test --- .../google_genai/adapters/transformation.py | 62 +++++++-- .../google_genai/test_google_genai_adapter.py | 127 +++++++++++++++++- 2 files changed, 179 insertions(+), 10 deletions(-) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 9d3f990b1aa..58a52666d38 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -8,8 +8,10 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, + ChatCompletionImageObject, ChatCompletionRequest, ChatCompletionSystemMessage, + ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, ChatCompletionToolChoiceValues, ChatCompletionToolMessage, @@ -385,13 +387,36 @@ class GoogleGenAIAdapter: if role == "user": # Handle user messages with potential function responses - combined_text = "" + content_parts: List[ + Union[ChatCompletionTextObject, ChatCompletionImageObject] + ] = [] tool_messages: List[ChatCompletionToolMessage] = [] for part in parts: if isinstance(part, dict): if "text" in part: - combined_text += part["text"] + content_parts.append( + cast( + ChatCompletionTextObject, + {"type": "text", "text": part["text"]}, + ) + ) + elif "inline_data" in part: + # Handle Base64 image data + inline_data = part["inline_data"] + mime_type = inline_data.get("mime_type", "image/jpeg") + data = inline_data.get("data", "") + content_parts.append( + cast( + ChatCompletionImageObject, + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{data}" + }, + }, + ) + ) elif "functionResponse" in part: # Transform function response to tool message func_response = part["functionResponse"] @@ -402,13 +427,33 @@ class GoogleGenAIAdapter: ) tool_messages.append(tool_message) elif isinstance(part, str): - combined_text += part + content_parts.append( + cast( + ChatCompletionTextObject, {"type": "text", "text": part} + ) + ) - # Add user message if there's text content - if combined_text: - messages.append( - ChatCompletionUserMessage(role="user", content=combined_text) - ) + # Add user message if there's content + if content_parts: + # If only one text part, use simple string format for backward compatibility + if ( + len(content_parts) == 1 + and isinstance(content_parts[0], dict) + and content_parts[0].get("type") == "text" + ): + text_part = cast(ChatCompletionTextObject, content_parts[0]) + messages.append( + ChatCompletionUserMessage( + role="user", content=text_part["text"] + ) + ) + else: + # Use multimodal format (array of content parts) + messages.append( + ChatCompletionUserMessage( + role="user", content=content_parts + ) + ) # Add tool messages messages.extend(tool_messages) @@ -468,7 +513,6 @@ class GoogleGenAIAdapter: Dict in Google GenAI generate_content response format """ - # Extract the main response content choice = response.choices[0] if response.choices else None if not choice: diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index e8882a1acb3..135881ad209 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1197,6 +1197,131 @@ async def test_agenerate_content_x_goog_api_key_header(): # Verify other expected headers assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" - + print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") print(f"✓ All headers: {list(headers.keys())}") + + +def test_inline_data_base64_image_transformation(): + """Test transformation of Gemini inline_data (Base64 images) to OpenAI format""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with Base64 image + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + {"text": "What's in this image?"}, + { + "inline_data": { + "mime_type": "image/jpeg", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 2, "Should have 2 content parts (text + image)" + + # Verify text part + text_part = content[0] + assert text_part["type"] == "text" + assert text_part["text"] == "What's in this image?" + + # Verify image part + image_part = content[1] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" in image_part["image_url"]["url"] + + +def test_inline_data_image_only_transformation(): + """Test transformation of Gemini inline_data with only image (no text)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only Base64 image (no text) + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + { + "inline_data": { + "mime_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 1, "Should have 1 content part (image only)" + + # Verify image part + image_part = content[0] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/png;base64,") + + +def test_inline_data_backward_compatibility_text_only(): + """Test that pure text messages still use simple string format (backward compatibility)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only text (no images) + model = "gpt-3.5-turbo" + contents = { + "role": "user", + "parts": [{"text": "Hello, how are you?"}] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is a simple string (not an array) for backward compatibility + content = completion_request["messages"][0]["content"] + assert isinstance(content, str), "Content should be a string for text-only messages (backward compatibility)" + assert content == "Hello, how are you?" From 9b1c5f7e360e7b448d9f16500e4ababba6f47b42 Mon Sep 17 00:00:00 2001 From: cantalupo555 Date: Sat, 3 Jan 2026 16:14:19 -0300 Subject: [PATCH 236/330] feat(zai): Add GLM-4.7 model with reasoning support (#18476) Add support for Z.AI GLM-4.7, latest flagship model with enhanced reasoning capabilities. Changes: - Add zai/glm-4.7 to model pricing with /bin/bash.60/M input, .20/M output - Add cached input pricing (/bin/bash.11/M) for GLM-4.7 - Add supports_reasoning flag to enable thinking parameter - Update ZAIChatConfig to support thinking parameter for models with reasoning - Update documentation with GLM-4.7 as latest flagship model - Add cached input column to pricing table (GLM-4.7 only) - Add tests for GLM-4.7 reasoning support and cost calculation - Update all examples to use GLM-4.7 Model specifications: - Context: 200K input, 128K output - Supports: reasoning, function calling, tool choice, prompt caching - Pricing: Same as GLM-4.6 with cache support See: https://docs.z.ai/guides/llm/glm-4.7 --- docs/my-website/docs/providers/zai.md | 36 +++++++++--------- litellm/llms/zai/chat/transformation.py | 11 +++++- ...odel_prices_and_context_window_backup.json | 14 +++++++ model_prices_and_context_window.json | 14 +++++++ .../llms/zai/test_zai_provider.py | 37 +++++++++++++++++++ 5 files changed, 94 insertions(+), 18 deletions(-) diff --git a/docs/my-website/docs/providers/zai.md b/docs/my-website/docs/providers/zai.md index 5055d0c1cdd..937ccd67680 100644 --- a/docs/my-website/docs/providers/zai.md +++ b/docs/my-website/docs/providers/zai.md @@ -19,7 +19,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -34,7 +34,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -51,7 +51,8 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet | Model Name | Function Call | Notes | |------------|---------------|-------| -| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | Latest flagship model, 200K context | +| glm-4.7 | `completion(model="zai/glm-4.7", messages)` | **Latest flagship**, 200K context, **Reasoning** | +| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | 200K context | | glm-4.5 | `completion(model="zai/glm-4.5", messages)` | 128K context | | glm-4.5v | `completion(model="zai/glm-4.5v", messages)` | Vision model | | glm-4.5-x | `completion(model="zai/glm-4.5-x", messages)` | Premium tier | @@ -62,16 +63,17 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet ## Model Pricing -| Model | Input ($/1M tokens) | Output ($/1M tokens) | Context Window | -|-------|---------------------|----------------------|----------------| -| glm-4.6 | $0.60 | $2.20 | 200K | -| glm-4.5 | $0.60 | $2.20 | 128K | -| glm-4.5v | $0.60 | $1.80 | 128K | -| glm-4.5-x | $2.20 | $8.90 | 128K | -| glm-4.5-air | $0.20 | $1.10 | 128K | -| glm-4.5-airx | $1.10 | $4.50 | 128K | -| glm-4-32b-0414-128k | $0.10 | $0.10 | 128K | -| glm-4.5-flash | **FREE** | **FREE** | 128K | +| Model | Input ($/1M tokens) | Output ($/1M tokens) | Cached Input ($/1M tokens) | Context Window | +|-------|---------------------|----------------------|---------------------------|----------------| +| glm-4.7 | $0.60 | $2.20 | $0.11 | 200K | +| glm-4.6 | $0.60 | $2.20 | - | 200K | +| glm-4.5 | $0.60 | $2.20 | - | 128K | +| glm-4.5v | $0.60 | $1.80 | - | 128K | +| glm-4.5-x | $2.20 | $8.90 | - | 128K | +| glm-4.5-air | $0.20 | $1.10 | - | 128K | +| glm-4.5-airx | $1.10 | $4.50 | - | 128K | +| glm-4-32b-0414-128k | $0.10 | $0.10 | - | 128K | +| glm-4.5-flash | **FREE** | **FREE** | - | 128K | ## Using with LiteLLM Proxy @@ -84,7 +86,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[{"role": "user", "content": "Hello, how are you?"}], ) @@ -98,9 +100,9 @@ print(response.choices[0].message.content) ```yaml model_list: - - model_name: glm-4.6 + - model_name: glm-4.7 litellm_params: - model: zai/glm-4.6 + model: zai/glm-4.7 api_key: os.environ/ZAI_API_KEY - model_name: glm-4.5-flash # Free tier litellm_params: @@ -121,7 +123,7 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ -H 'Content-Type: application/json' \ -H 'Authorization: Bearer sk-1234' \ -d '{ - "model": "glm-4.6", + "model": "glm-4.7", "messages": [ { "role": "user", diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index 47b314d4e0d..4380256f0a4 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -20,7 +20,7 @@ class ZAIChatConfig(OpenAIGPTConfig): return api_base, dynamic_api_key def get_supported_openai_params(self, model: str) -> list: - return [ + base_params = [ "max_tokens", "stream", "stream_options", @@ -31,3 +31,12 @@ class ZAIChatConfig(OpenAIGPTConfig): "tool_choice", ] + import litellm + + try: + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): + base_params.append("thinking") + except Exception: + pass + + return base_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c585dac9063..d32adf54b5e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29649,6 +29649,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index df286e6540a..81b4469f24c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29691,6 +29691,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index a3d47d666bc..d1e4359d048 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -1,6 +1,7 @@ """ Tests for Z.AI (Zhipu AI) provider - GLM models """ + import json import math @@ -50,10 +51,12 @@ def test_zai_in_provider_lists(): def test_zai_models_in_model_cost(): """Test that ZAI models are in the model cost map""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") zai_models = [ + "zai/glm-4.7", "zai/glm-4.6", "zai/glm-4.5", "zai/glm-4.5v", @@ -72,6 +75,7 @@ def test_zai_models_in_model_cost(): def test_zai_glm46_cost_calculation(): """Test the cost calculation for glm-4.6""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -92,6 +96,7 @@ def test_zai_glm46_cost_calculation(): def test_zai_flash_model_is_free(): """Test that glm-4.5-flash has zero cost""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -102,6 +107,38 @@ def test_zai_flash_model_is_free(): assert info["output_cost_per_token"] == 0 +def test_glm47_supports_reasoning(): + """Test that GLM-4.7 supports reasoning""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + key = "zai/glm-4.7" + assert key in litellm.model_cost, f"Model {key} not found in model_cost" + + info = litellm.model_cost[key] + assert info["supports_reasoning"] is True + + +def test_glm47_cost_calculation(): + """Test cost calculation for GLM-4.7""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prompt_cost, completion_cost = cost_per_token( + model="zai/glm-4.7", + prompt_tokens=1000000, # 1M tokens + completion_tokens=1000000, + ) + + # GLM-4.7: $0.6/M input, $2.2/M output (same as GLM-4.6) + assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6) + assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) + + @pytest.mark.asyncio async def test_zai_completion_call(respx_mock, zai_response, monkeypatch): """Test completion call with zai provider using mocked response""" From 89b4a6d67c2e603e294bcf2f35bf34bd7e9ba2ae Mon Sep 17 00:00:00 2001 From: Anders Kaseorg Date: Sat, 3 Jan 2026 11:15:15 -0800 Subject: [PATCH 237/330] Allow installation with current grpcio on old Python (#18473) Instead of limiting grpcio < 1.68.0, specifically exclude the versions affected by the reconnect bug, and allow installation with either older or newer versions. Signed-off-by: Anders Kaseorg --- poetry.lock | 72 +----------------------------------------------- pyproject.toml | 2 +- requirements.txt | 2 +- 3 files changed, 3 insertions(+), 73 deletions(-) diff --git a/poetry.lock b/poetry.lock index ee97c00594c..a0a0f8540e5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2273,75 +2273,6 @@ googleapis-common-protos = {version = ">=1.56.0,<2.0.0", extras = ["grpc"]} grpcio = ">=1.44.0,<2.0.0" protobuf = ">=3.20.2,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" -[[package]] -name = "grpcio" -version = "1.67.1" -description = "HTTP/2-based RPC framework" -optional = false -python-versions = ">=3.8" -groups = ["main", "dev", "proxy-dev"] -markers = "python_version < \"3.14\"" -files = [ - {file = "grpcio-1.67.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:8b0341d66a57f8a3119b77ab32207072be60c9bf79760fa609c5609f2deb1f3f"}, - {file = "grpcio-1.67.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:f5a27dddefe0e2357d3e617b9079b4bfdc91341a91565111a21ed6ebbc51b22d"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:43112046864317498a33bdc4797ae6a268c36345a910de9b9c17159d8346602f"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c9b929f13677b10f63124c1a410994a401cdd85214ad83ab67cc077fc7e480f0"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e7d1797a8a3845437d327145959a2c0c47c05947c9eef5ff1a4c80e499dcc6fa"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:0489063974d1452436139501bf6b180f63d4977223ee87488fe36858c5725292"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:9fd042de4a82e3e7aca44008ee2fb5da01b3e5adb316348c21980f7f58adc311"}, - {file = "grpcio-1.67.1-cp310-cp310-win32.whl", hash = "sha256:638354e698fd0c6c76b04540a850bf1db27b4d2515a19fcd5cf645c48d3eb1ed"}, - {file = "grpcio-1.67.1-cp310-cp310-win_amd64.whl", hash = "sha256:608d87d1bdabf9e2868b12338cd38a79969eaf920c89d698ead08f48de9c0f9e"}, - {file = "grpcio-1.67.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:7818c0454027ae3384235a65210bbf5464bd715450e30a3d40385453a85a70cb"}, - {file = "grpcio-1.67.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ea33986b70f83844cd00814cee4451055cd8cab36f00ac64a31f5bb09b31919e"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:c7a01337407dd89005527623a4a72c5c8e2894d22bead0895306b23c6695698f"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:80b866f73224b0634f4312a4674c1be21b2b4afa73cb20953cbbb73a6b36c3cc"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fff78ba10d4250bfc07a01bd6254a6d87dc67f9627adece85c0b2ed754fa96"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:8a23cbcc5bb11ea7dc6163078be36c065db68d915c24f5faa4f872c573bb400f"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:1a65b503d008f066e994f34f456e0647e5ceb34cfcec5ad180b1b44020ad4970"}, - {file = "grpcio-1.67.1-cp311-cp311-win32.whl", hash = "sha256:e29ca27bec8e163dca0c98084040edec3bc49afd10f18b412f483cc68c712744"}, - {file = "grpcio-1.67.1-cp311-cp311-win_amd64.whl", hash = "sha256:786a5b18544622bfb1e25cc08402bd44ea83edfb04b93798d85dca4d1a0b5be5"}, - {file = "grpcio-1.67.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:267d1745894200e4c604958da5f856da6293f063327cb049a51fe67348e4f953"}, - {file = "grpcio-1.67.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:85f69fdc1d28ce7cff8de3f9c67db2b0ca9ba4449644488c1e0303c146135ddb"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:f26b0b547eb8d00e195274cdfc63ce64c8fc2d3e2d00b12bf468ece41a0423a0"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4422581cdc628f77302270ff839a44f4c24fdc57887dc2a45b7e53d8fc2376af"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1d7616d2ded471231c701489190379e0c311ee0a6c756f3c03e6a62b95a7146e"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8a00efecde9d6fcc3ab00c13f816313c040a28450e5e25739c24f432fc6d3c75"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:699e964923b70f3101393710793289e42845791ea07565654ada0969522d0a38"}, - {file = "grpcio-1.67.1-cp312-cp312-win32.whl", hash = "sha256:4e7b904484a634a0fff132958dabdb10d63e0927398273917da3ee103e8d1f78"}, - {file = "grpcio-1.67.1-cp312-cp312-win_amd64.whl", hash = "sha256:5721e66a594a6c4204458004852719b38f3d5522082be9061d6510b455c90afc"}, - {file = "grpcio-1.67.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:aa0162e56fd10a5547fac8774c4899fc3e18c1aa4a4759d0ce2cd00d3696ea6b"}, - {file = "grpcio-1.67.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:beee96c8c0b1a75d556fe57b92b58b4347c77a65781ee2ac749d550f2a365dc1"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:a93deda571a1bf94ec1f6fcda2872dad3ae538700d94dc283c672a3b508ba3af"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e6f255980afef598a9e64a24efce87b625e3e3c80a45162d111a461a9f92955"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e838cad2176ebd5d4a8bb03955138d6589ce9e2ce5d51c3ada34396dbd2dba8"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:a6703916c43b1d468d0756c8077b12017a9fcb6a1ef13faf49e67d20d7ebda62"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:917e8d8994eed1d86b907ba2a61b9f0aef27a2155bca6cbb322430fc7135b7bb"}, - {file = "grpcio-1.67.1-cp313-cp313-win32.whl", hash = "sha256:e279330bef1744040db8fc432becc8a727b84f456ab62b744d3fdb83f327e121"}, - {file = "grpcio-1.67.1-cp313-cp313-win_amd64.whl", hash = "sha256:fa0c739ad8b1996bd24823950e3cb5152ae91fca1c09cc791190bf1627ffefba"}, - {file = "grpcio-1.67.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:178f5db771c4f9a9facb2ab37a434c46cb9be1a75e820f187ee3d1e7805c4f65"}, - {file = "grpcio-1.67.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f3e49c738396e93b7ba9016e153eb09e0778e776df6090c1b8c91877cc1c426"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:24e8a26dbfc5274d7474c27759b54486b8de23c709d76695237515bc8b5baeab"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3b6c16489326d79ead41689c4b84bc40d522c9a7617219f4ad94bc7f448c5085"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60e6a4dcf5af7bbc36fd9f81c9f372e8ae580870a9e4b6eafe948cd334b81cf3"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:95b5f2b857856ed78d72da93cd7d09b6db8ef30102e5e7fe0961fe4d9f7d48e8"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b49359977c6ec9f5d0573ea4e0071ad278ef905aa74e420acc73fd28ce39e9ce"}, - {file = "grpcio-1.67.1-cp38-cp38-win32.whl", hash = "sha256:f5b76ff64aaac53fede0cc93abf57894ab2a7362986ba22243d06218b93efe46"}, - {file = "grpcio-1.67.1-cp38-cp38-win_amd64.whl", hash = "sha256:804c6457c3cd3ec04fe6006c739579b8d35c86ae3298ffca8de57b493524b771"}, - {file = "grpcio-1.67.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:a25bdea92b13ff4d7790962190bf6bf5c4639876e01c0f3dda70fc2769616335"}, - {file = "grpcio-1.67.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:cdc491ae35a13535fd9196acb5afe1af37c8237df2e54427be3eecda3653127e"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:85f862069b86a305497e74d0dc43c02de3d1d184fc2c180993aa8aa86fbd19b8"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ec74ef02010186185de82cc594058a3ccd8d86821842bbac9873fd4a2cf8be8d"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01f616a964e540638af5130469451cf580ba8c7329f45ca998ab66e0c7dcdb04"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:299b3d8c4f790c6bcca485f9963b4846dd92cf6f1b65d3697145d005c80f9fe8"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:60336bff760fbb47d7e86165408126f1dded184448e9a4c892189eb7c9d3f90f"}, - {file = "grpcio-1.67.1-cp39-cp39-win32.whl", hash = "sha256:5ed601c4c6008429e3d247ddb367fe8c7259c355757448d7c1ef7bd4a6739e8e"}, - {file = "grpcio-1.67.1-cp39-cp39-win_amd64.whl", hash = "sha256:5db70d32d6703b89912af16d6d45d78406374a8b8ef0d28140351dd0ec610e98"}, - {file = "grpcio-1.67.1.tar.gz", hash = "sha256:3dc2ed4cabea4dc14d5e708c2b426205956077cc5de419b4d4079315017e9732"}, -] - -[package.extras] -protobuf = ["grpcio-tools (>=1.67.1)"] - [[package]] name = "grpcio" version = "1.76.0" @@ -2349,7 +2280,6 @@ description = "HTTP/2-based RPC framework" optional = false python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] -markers = "python_version >= \"3.14\"" files = [ {file = "grpcio-1.76.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:65a20de41e85648e00305c1bb09a3598f840422e522277641145a32d42dcefcc"}, {file = "grpcio-1.76.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:40ad3afe81676fd9ec6d9d406eda00933f218038433980aa19d401490e46ecde"}, @@ -8051,4 +7981,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "b010d9da7f5a765670932b78d720aae4fcb819daba050683ee125b4367972419" +content-hash = "7eed2b2c25173a275ac83c55fd901b9b84663b1d7daa54f0e78b30bf1c8f0e3e" diff --git a/pyproject.toml b/pyproject.toml index f929fb94cb0..3b09119a748 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ soundfile = {version = "^0.12.1", optional = true} # - 1.68.0-1.68.1 has reconnect bug (https://github.com/grpc/grpc/issues/38290) # - 1.75.0+ has Python 3.14 wheels and bug fix grpcio = [ - {version = ">=1.62.3,<1.68.0", python = "<3.14"}, + {version = ">=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0", python = "<3.14"}, {version = ">=1.75.0", python = ">=3.14"}, ] diff --git a/requirements.txt b/requirements.txt index 3bc968c8cb8..06a7c17336c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -41,7 +41,7 @@ opentelemetry-api==1.25.0 opentelemetry-sdk==1.25.0 opentelemetry-exporter-otlp==1.25.0 # grpcio: 1.68.0-1.68.1 has reconnect bug (#38290), 1.75+ has Python 3.14 wheels + fix -grpcio>=1.62.3,<1.68.0; python_version < "3.14" +grpcio>=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0; python_version < "3.14" grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests From 099e108b51df00d2c8e855b459264973b1bb03b0 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:15:45 +0800 Subject: [PATCH 238/330] fix: correctly route codestral chat and FIM endpoints (#18467) Fixed duplicate condition that made text-completion-codestral provider unreachable. Now: - codestral.mistral.ai/v1/chat/completions -> codestral - codestral.mistral.ai/v1/fim/completions -> text-completion-codestral Fixes #18464 Signed-off-by: majiayu000 <1835304752@qq.com> --- .../get_llm_provider_logic.py | 4 +- .../test_codestral_provider_routing.py | 69 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 164e2a73e65..b753e9fa8b5 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -229,10 +229,10 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://api.ai21.com/studio/v1": custom_llm_provider = "ai21_chat" dynamic_api_key = get_secret_str("AI21_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/chat/completions": custom_llm_provider = "codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/fim/completions": custom_llm_provider = "text-completion-codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") elif endpoint == "app.empower.dev/api/v1": diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py new file mode 100644 index 00000000000..1a6ed51afd0 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py @@ -0,0 +1,69 @@ +""" +Unit tests for codestral provider routing. + +These tests verify that the chat and FIM endpoints for codestral +are correctly routed to different providers: +- Chat endpoint -> codestral provider +- FIM endpoint -> text-completion-codestral provider + +Related issue: https://github.com/BerriAI/litellm/issues/18464 +""" +import pytest + +import litellm + + +class TestCodestralProviderRouting: + """Tests for codestral endpoint routing in get_llm_provider""" + + def test_codestral_chat_endpoint_routes_to_codestral_provider(self): + """ + Test that the codestral chat endpoint routes to the 'codestral' provider. + + The chat/completions endpoint should be handled by the codestral provider. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + assert custom_llm_provider == "codestral" + + def test_codestral_fim_endpoint_routes_to_text_completion_provider(self): + """ + Test that the codestral FIM endpoint routes to 'text-completion-codestral'. + + The fim/completions endpoint should be handled by the + text-completion-codestral provider for fill-in-the-middle completions. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert custom_llm_provider == "text-completion-codestral" + + def test_codestral_endpoints_are_different_providers(self): + """ + Test that chat and FIM endpoints route to different providers. + + This is the core fix for issue #18464 - previously both endpoints + would route to 'codestral' due to duplicate conditions. + """ + _, chat_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + _, fim_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert chat_provider != fim_provider + assert chat_provider == "codestral" + assert fim_provider == "text-completion-codestral" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 64cfe75bfd398aa856c117aedfb6b72c70f48d16 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:17:38 +0800 Subject: [PATCH 239/330] fix: extract pure base64 data from data URLs for Ollama (#18465) Fix Ollama_chatException "illegal base64 data at input byte 4" error when using images with ollama_chat provider. Ollama expects pure base64 data, not the full data URL format (data:image/png;base64,...). Fixes #18338 Signed-off-by: majiayu000 <1835304752@qq.com> --- .../prompt_templates/common_utils.py | 32 +++- .../test_extract_base64_image.py | 156 ++++++++++++++++++ 2 files changed, 185 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_extract_base64_image.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ca2a092dbc8..b100b9b516b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1087,9 +1087,35 @@ def _parse_content_for_reasoning( return None, message_text +def _extract_base64_data(image_url: str) -> str: + """ + Extract pure base64 data from an image URL. + + If the URL is a data URL (e.g., "data:image/png;base64,iVBOR..."), + extract and return only the base64 data portion. + Otherwise, return the original URL unchanged. + + This is needed for providers like Ollama that expect pure base64 data + rather than full data URLs. + + Args: + image_url: The image URL or data URL to process + + Returns: + The base64 data if it's a data URL, otherwise the original URL + """ + if image_url.startswith("data:") and ";base64," in image_url: + return image_url.split(";base64,", 1)[1] + return image_url + + def extract_images_from_message(message: AllMessageValues) -> List[str]: """ - Extract images from a message + Extract images from a message. + + For data URLs (e.g., "data:image/png;base64,iVBOR..."), only the base64 + data portion is extracted. This is required for providers like Ollama + that expect pure base64 data rather than full data URLs. """ images = [] message_content = message.get("content") @@ -1098,7 +1124,7 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]: image_url = m.get("image_url") if image_url: if isinstance(image_url, str): - images.append(image_url) + images.append(_extract_base64_data(image_url)) elif isinstance(image_url, dict) and "url" in image_url: - images.append(image_url["url"]) + images.append(_extract_base64_data(image_url["url"])) return images diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py new file mode 100644 index 00000000000..b17c02d7006 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py @@ -0,0 +1,156 @@ +""" +Unit tests for _extract_base64_data and extract_images_from_message functions. + +These tests verify that base64 image data is correctly extracted from data URLs, +which fixes the Ollama error "illegal base64 data at input byte 4". + +Related issue: https://github.com/BerriAI/litellm/issues/18338 +""" +import pytest + +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_base64_data, + extract_images_from_message, +) + + +class TestExtractBase64Data: + """Tests for _extract_base64_data function""" + + def test_extract_base64_from_png_data_url(self): + """Test extracting base64 data from a PNG data URL""" + data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + expected = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_jpeg_data_url(self): + """Test extracting base64 data from a JPEG data URL""" + data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD" + expected = "/9j/4AAQSkZJRgABAQAAAQABAAD" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_gif_data_url(self): + """Test extracting base64 data from a GIF data URL""" + data_url = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP" + expected = "R0lGODlhAQABAIAAAAAAAP" + assert _extract_base64_data(data_url) == expected + + def test_regular_url_unchanged(self): + """Test that regular HTTP URLs are returned unchanged""" + url = "https://example.com/image.png" + assert _extract_base64_data(url) == url + + def test_file_path_unchanged(self): + """Test that file paths are returned unchanged""" + path = "/path/to/image.png" + assert _extract_base64_data(path) == path + + def test_data_url_without_base64_unchanged(self): + """Test that data URLs without base64 encoding are returned unchanged""" + # This is a data URL with URL encoding, not base64 + url = "data:text/plain,Hello%20World" + assert _extract_base64_data(url) == url + + def test_base64_data_with_special_chars(self): + """Test extracting base64 data that contains valid special characters""" + # Base64 can contain +, /, and = characters + data_url = "data:image/png;base64,abc+def/ghi===" + expected = "abc+def/ghi===" + assert _extract_base64_data(data_url) == expected + + +class TestExtractImagesFromMessage: + """Tests for extract_images_from_message function""" + + def test_extract_from_message_with_data_url_string(self): + """Test extracting images when image_url is a string data URL""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_data_url_dict(self): + """Test extracting images when image_url is a dict with url key""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_regular_url(self): + """Test that regular URLs are preserved""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["https://example.com/image.png"] + + def test_extract_multiple_images(self): + """Test extracting multiple images from a single message""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,image1base64", + }, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,image2base64"}, + }, + { + "type": "image_url", + "image_url": "https://example.com/image3.png", + }, + ], + } + result = extract_images_from_message(message) + assert result == [ + "image1base64", + "image2base64", + "https://example.com/image3.png", + ] + + def test_empty_content(self): + """Test message with empty content""" + message = {"role": "user", "content": []} + result = extract_images_from_message(message) + assert result == [] + + def test_no_images_in_content(self): + """Test message with content but no images""" + message = { + "role": "user", + "content": [{"type": "text", "text": "Hello world"}], + } + result = extract_images_from_message(message) + assert result == [] + + def test_string_content(self): + """Test message with string content (no images possible)""" + message = {"role": "user", "content": "Hello world"} + result = extract_images_from_message(message) + assert result == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 7d81d245fb2665f279c69da1a2066532b07f3045 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:18:09 +0800 Subject: [PATCH 240/330] fix: align prometheus metric names with DEFINED_PROMETHEUS_METRICS (#18463) Fix metric name inconsistency for litellm_remaining_requests_metric and litellm_remaining_tokens_metric. The factory received names without the _metric suffix, causing _is_metric_enabled to fail when users configured these metrics in prometheus_metrics_config. Fixes #18221 Signed-off-by: majiayu000 <1835304752@qq.com> --- litellm/integrations/prometheus.py | 4 +- ...test_prometheus_metric_name_consistency.py | 106 ++++++++++++++++++ 2 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 20f1357a1c8..c01f7481277 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -214,7 +214,7 @@ class PrometheusLogger(CustomLogger): # Remaining Rate Limit for model self.litellm_remaining_requests_metric = self._gauge_factory( - "litellm_remaining_requests", + "litellm_remaining_requests_metric", "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_requests_metric" @@ -222,7 +222,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_tokens_metric = self._gauge_factory( - "litellm_remaining_tokens", + "litellm_remaining_tokens_metric", "remaining tokens for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_tokens_metric" diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py new file mode 100644 index 00000000000..9658eff3cc5 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py @@ -0,0 +1,106 @@ +""" +Unit tests for prometheus metric name consistency + +This test ensures that the metric names used when creating Prometheus metrics +match the names defined in DEFINED_PROMETHEUS_METRICS, so that metric filtering +configuration works correctly. + +Related issue: https://github.com/BerriAI/litellm/issues/18221 +""" +from typing import get_args + +import pytest + + +def test_remaining_requests_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_requests_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_requests_metric" in defined_metrics + ), "litellm_remaining_requests_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_remaining_tokens_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_tokens_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_tokens_metric" in defined_metrics + ), "litellm_remaining_tokens_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_prometheus_metric_labels_have_remaining_metrics(): + """ + Test that PrometheusMetricLabels has label definitions for remaining metrics. + + This ensures that the labels can be retrieved when creating the metrics. + """ + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Test that labels can be retrieved for remaining metrics + remaining_requests_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_requests_metric" + ) + remaining_tokens_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_tokens_metric" + ) + + assert isinstance( + remaining_requests_labels, list + ), "Labels for litellm_remaining_requests_metric should be a list" + assert isinstance( + remaining_tokens_labels, list + ), "Labels for litellm_remaining_tokens_metric should be a list" + + # These metrics should have api_provider and api_base labels + assert ( + "api_provider" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_provider label" + assert ( + "api_base" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_base label" + assert ( + "api_provider" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_provider label" + assert ( + "api_base" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_base label" + + +def test_all_defined_metrics_have_consistent_naming(): + """ + Test that all metrics defined in DEFINED_PROMETHEUS_METRICS follow + a consistent naming convention. + + This helps prevent similar inconsistencies in the future. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + for metric_name in defined_metrics: + # All metrics should start with 'litellm_' + assert metric_name.startswith( + "litellm_" + ), f"Metric {metric_name} should start with 'litellm_'" + + +if __name__ == "__main__": + test_remaining_requests_metric_name_in_defined_metrics() + test_remaining_tokens_metric_name_in_defined_metrics() + test_prometheus_metric_labels_have_remaining_metrics() + test_all_defined_metrics_have_consistent_naming() + print("All prometheus metric name consistency tests passed!") From 1452f0150551193bd26227a348f2e0acd6d294fd Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 12:12:24 -0800 Subject: [PATCH 241/330] refactor: lazy load get_llm_provider and remove_index_from_tool_calls (#18608) --- litellm/_lazy_imports_registry.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 1a54e77998a..93fa8b39af2 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -32,6 +32,7 @@ UTILS_NAMES = ( "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", "ModelResponseListIterator", "get_valid_models", "timeout", + "get_llm_provider", "remove_index_from_tool_calls", ) # Token counter names that support lazy loading via _lazy_import_token_counter @@ -336,6 +337,8 @@ _UTILS_IMPORT_MAP = { "ModelResponseListIterator": (".utils", "ModelResponseListIterator"), "get_valid_models": (".utils", "get_valid_models"), "timeout": (".timeout", "timeout"), + "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"), + "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"), } _COST_CALCULATOR_IMPORT_MAP = { From 4904ed394eeace50bf951c8a0c9b2fd49b77b2c1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 12:36:26 -0800 Subject: [PATCH 242/330] Edit path for SSO Settings --- .../hooks/sso/useEditSSOSettings.ts | 38 +++ .../(dashboard)/hooks/sso/useSSOSettings.ts | 9 +- .../Modals/AddSSOSettingsModal.test.tsx | 2 +- .../Modals/AddSSOSettingsModal.tsx | 245 +++------------- .../Modals/BaseSSOSettingsForm.tsx | 249 ++++++++++++++++ .../Modals/EditSSOSettingsModal.tsx | 131 +++++++++ .../SSOSettings/RedactableField.tsx | 38 +++ .../AdminSettings/SSOSettings/SSOSettings.tsx | 84 ++++-- .../AdminSettings/SSOSettings/constants.ts | 15 + .../AdminSettings/SSOSettings/utils.test.ts | 274 ++++++++++++++++++ .../AdminSettings/SSOSettings/utils.ts | 54 ++++ 11 files changed, 894 insertions(+), 245 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts new file mode 100644 index 00000000000..69e52d0ff25 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts @@ -0,0 +1,38 @@ +import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { updateSSOSettings } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export interface EditSSOSettingsParams { + google_client_id?: string | null; + google_client_secret?: string | null; + microsoft_client_id?: string | null; + microsoft_client_secret?: string | null; + microsoft_tenant?: string | null; + generic_client_id?: string | null; + generic_client_secret?: string | null; + generic_authorization_endpoint?: string | null; + generic_token_endpoint?: string | null; + generic_userinfo_endpoint?: string | null; + proxy_base_url?: string | null; + user_email?: string | null; + sso_provider?: string | null; + role_mappings?: any; + [key: string]: any; +} + +export interface EditSSOSettingsResponse { + [key: string]: any; +} + +export const useEditSSOSettings = (): UseMutationResult => { + const { accessToken } = useAuthorized(); + + return useMutation({ + mutationFn: async (params: EditSSOSettingsParams) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return await updateSSOSettings(accessToken, params); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts index 6453b35ff93..3e09c3c2ca8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts @@ -27,7 +27,14 @@ export interface SSOSettingsValues { proxy_base_url: string | null; user_email: string | null; ui_access_mode: string | null; - role_mappings: string | null; + role_mappings: { + provider: string; + group_claim: string; + default_role: "internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer"; + roles: { + [key: string]: string[]; + }; + }; } export interface SSOSettingsResponse { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx index 13363a11643..aae28191031 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -17,7 +17,7 @@ describe("AddSSOSettingsModal", () => { const onCancel = vi.fn(); const onSuccess = vi.fn(); - render(); + render(); expect(screen.getByText("SSO Provider")).toBeInTheDocument(); expect(screen.getByText("Cancel")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx index a4ef2938f60..7af6240b19e 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx @@ -1,145 +1,36 @@ "use client"; import NotificationsManager from "@/components/molecules/notifications_manager"; -import { updateSSOSettings } from "@/components/networking"; import { parseErrorMessage } from "@/components/shared/errorUtils"; -import { TextInput } from "@tremor/react"; -import { Button as Button2, Form, Input, Modal, Select } from "antd"; -import React, { useState } from "react"; +import { Button, Form, Modal, Space } from "antd"; +import React from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import { processSSOSettingsPayload } from "../utils"; interface AddSSOSettingsModalProps { isVisible: boolean; onCancel: () => void; onSuccess: () => void; - accessToken: string | null; } -const ssoProviderLogoMap: Record = { - google: "https://artificialanalysis.ai/img/logos/google_small.svg", - microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", - okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png", - generic: "", -}; - -// Define the SSO provider configuration type -interface SSOProviderConfig { - envVarMap: Record; - fields: Array<{ - label: string; - name: string; - placeholder?: string; - }>; -} - -// Define configurations for each SSO provider -const ssoProviderConfigs: Record = { - google: { - envVarMap: { - google_client_id: "GOOGLE_CLIENT_ID", - google_client_secret: "GOOGLE_CLIENT_SECRET", - }, - fields: [ - { label: "Google Client ID", name: "google_client_id" }, - { label: "Google Client Secret", name: "google_client_secret" }, - ], - }, - microsoft: { - envVarMap: { - microsoft_client_id: "MICROSOFT_CLIENT_ID", - microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", - microsoft_tenant: "MICROSOFT_TENANT", - }, - fields: [ - { label: "Microsoft Client ID", name: "microsoft_client_id" }, - { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, - { label: "Microsoft Tenant", name: "microsoft_tenant" }, - ], - }, - okta: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { - label: "Authorization Endpoint", - name: "generic_authorization_endpoint", - placeholder: "https://your-domain/authorize", - }, - { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, - { - label: "Userinfo Endpoint", - name: "generic_userinfo_endpoint", - placeholder: "https://your-domain/userinfo", - }, - ], - }, - generic: { - envVarMap: { - generic_client_id: "GENERIC_CLIENT_ID", - generic_client_secret: "GENERIC_CLIENT_SECRET", - generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", - generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", - generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", - }, - fields: [ - { label: "Generic Client ID", name: "generic_client_id" }, - { label: "Generic Client Secret", name: "generic_client_secret" }, - { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, - { label: "Token Endpoint", name: "generic_token_endpoint" }, - { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, - ], - }, -}; - -const AddSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess, accessToken }) => { +const AddSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { const [form] = Form.useForm(); - const [isSubmitting, setIsSubmitting] = useState(false); + const { mutateAsync, isPending } = useEditSSOSettings(); // Enhanced form submission handler const handleFormSubmit = async (formValues: Record) => { - if (!accessToken) { - NotificationsManager.fromBackend("No access token available"); - return; - } + const payload = processSSOSettingsPayload(formValues); - setIsSubmitting(true); - try { - // Save SSO settings using the new API - await updateSSOSettings(accessToken, formValues); - - NotificationsManager.success("SSO settings added successfully"); - - // Reset form and close modal - form.resetFields(); - onSuccess(); - } catch (error: unknown) { - NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); - } finally { - setIsSubmitting(false); - } - }; - - // Helper function to render provider fields - const renderProviderFields = (provider: string) => { - const config = ssoProviderConfigs[provider]; - if (!config) return null; - - return config.fields.map((field) => ( - - {field.name.includes("client") ? : } - - )); + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings added successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); }; const handleCancel = () => { @@ -148,91 +39,23 @@ const AddSSOSettingsModal: React.FC = ({ isVisible, on }; return ( - -
- - - - - prevValues.sso_provider !== currentValues.sso_provider} - > - {({ getFieldValue }) => { - const provider = getFieldValue("sso_provider"); - return provider ? renderProviderFields(provider) : null; - }} - - - - - - value?.trim()} - rules={[ - { required: true, message: "Please enter the proxy base url" }, - { - pattern: /^https?:\/\/.+/, - message: "URL must start with http:// or https://", - }, - { - validator: (_, value) => { - // Only check for trailing slash if the URL starts with http:// or https:// - if (value && /^https?:\/\/.+/.test(value) && value.endsWith("/")) { - return Promise.reject("URL must not end with a trailing slash"); - } - return Promise.resolve(); - }, - }, - ]} - > - - - -
- Cancel - - Add SSO - -
-
+ + + + + } + onCancel={handleCancel} + > + ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx new file mode 100644 index 00000000000..a4b36e5190e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -0,0 +1,249 @@ +"use client"; + +import { TextInput } from "@tremor/react"; +import { Checkbox, Form, Input, Select } from "antd"; +import React from "react"; +import { ssoProviderLogoMap, ssoProviderDisplayNames } from "../constants"; + +export interface BaseSSOSettingsFormProps { + form: any; // Replace with proper Form type if available + onFormSubmit: (formValues: Record) => Promise; +} + +// Define the SSO provider configuration type +export interface SSOProviderConfig { + envVarMap: Record; + fields: Array<{ + label: string; + name: string; + placeholder?: string; + }>; +} + +// Define configurations for each SSO provider +export const ssoProviderConfigs: Record = { + google: { + envVarMap: { + google_client_id: "GOOGLE_CLIENT_ID", + google_client_secret: "GOOGLE_CLIENT_SECRET", + }, + fields: [ + { label: "Google Client ID", name: "google_client_id" }, + { label: "Google Client Secret", name: "google_client_secret" }, + ], + }, + microsoft: { + envVarMap: { + microsoft_client_id: "MICROSOFT_CLIENT_ID", + microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", + microsoft_tenant: "MICROSOFT_TENANT", + }, + fields: [ + { label: "Microsoft Client ID", name: "microsoft_client_id" }, + { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, + { label: "Microsoft Tenant", name: "microsoft_tenant" }, + ], + }, + okta: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { + label: "Authorization Endpoint", + name: "generic_authorization_endpoint", + placeholder: "https://your-domain/authorize", + }, + { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, + { + label: "Userinfo Endpoint", + name: "generic_userinfo_endpoint", + placeholder: "https://your-domain/userinfo", + }, + ], + }, + generic: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, + { label: "Token Endpoint", name: "generic_token_endpoint" }, + { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, + ], + }, +}; + +// Helper function to render provider fields +export const renderProviderFields = (provider: string) => { + const config = ssoProviderConfigs[provider]; + if (!config) return null; + + return config.fields.map((field) => ( + + {field.name.includes("client") ? : } + + )); +}; + +const BaseSSOSettingsForm: React.FC = ({ form, onFormSubmit }) => { + return ( +
+
+ + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider ? renderProviderFields(provider) : null; + }} + + + + + + value?.trim()} + rules={[ + { required: true, message: "Please enter the proxy base url" }, + { + pattern: /^https?:\/\/.+/, + message: "URL must start with http:// or https://", + }, + { + validator: (_, value) => { + // Only check for trailing slash if the URL starts with http:// or https:// + if (value && /^https?:\/\/.+/.test(value) && value.endsWith("/")) { + return Promise.reject("URL must not end with a trailing slash"); + } + return Promise.resolve(); + }, + }, + ]} + > + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider === "okta" || provider === "generic" ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + <> + + + + + + + + + + + + + + + + + + + + + ) : null; + }} + +
+
+ ); +}; + +export default BaseSSOSettingsForm; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx new file mode 100644 index 00000000000..297698a7ba0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -0,0 +1,131 @@ +"use client"; + +import { Button, Form, Modal, Space } from "antd"; +import React, { useEffect } from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; + +interface EditSSOSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const EditSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + + // Use react-query hooks for SSO settings + const ssoSettings = useSSOSettings(); + const { mutateAsync, isPending } = useEditSSOSettings(); + useEffect(() => { + if (isVisible && ssoSettings.data && ssoSettings.data.values) { + const ssoData = ssoSettings.data; + console.log("Raw SSO data received:", ssoData); // Debug log + console.log("SSO values:", ssoData.values); // Debug log + console.log("user_email from API:", ssoData.values.user_email); // Debug log + + // Determine which SSO provider is configured + let selectedProvider = null; + if (ssoData.values.google_client_id) { + selectedProvider = "google"; + } else if (ssoData.values.microsoft_client_id) { + selectedProvider = "microsoft"; + } else if (ssoData.values.generic_client_id) { + // Check if it looks like Okta based on endpoints + if ( + ssoData.values.generic_authorization_endpoint?.includes("okta") || + ssoData.values.generic_authorization_endpoint?.includes("auth0") + ) { + selectedProvider = "okta"; + } else { + selectedProvider = "generic"; + } + } + + // Extract role mappings if they exist + let roleMappingFields = {}; + if (ssoData.values.role_mappings) { + const roleMappings = ssoData.values.role_mappings; + + // Helper function to join arrays into comma-separated strings + const joinTeams = (teams: string[] | undefined): string => { + if (!teams || teams.length === 0) return ""; + return teams.join(", "); + }; + + roleMappingFields = { + use_role_mappings: true, + group_claim: roleMappings.group_claim, + default_role: roleMappings.default_role || "internal_user", + proxy_admin_teams: joinTeams(roleMappings.roles?.proxy_admin), + admin_viewer_teams: joinTeams(roleMappings.roles?.proxy_admin_viewer), + internal_user_teams: joinTeams(roleMappings.roles?.internal_user), + internal_viewer_teams: joinTeams(roleMappings.roles?.internal_user_viewer), + }; + } + + // Set form values with existing data (excluding UI access control fields) + const formValues = { + sso_provider: selectedProvider, + ...ssoData.values, + ...roleMappingFields, + }; + + console.log("Setting form values:", formValues); // Debug log + + // Clear form first, then set values with a small delay to ensure proper initialization + form.resetFields(); + setTimeout(() => { + form.setFieldsValue(formValues); + console.log("Form values set, current form values:", form.getFieldsValue()); // Debug log + }, 100); + } + }, [isVisible, ssoSettings.data, form]); + + // Enhanced form submission handler + const handleFormSubmit = async (formValues: Record) => { + const payload = processSSOSettingsPayload(formValues); + + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings updated successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + + + + + } + onCancel={handleCancel} + > + + + ); +}; + +export default EditSSOSettingsModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx new file mode 100644 index 00000000000..44fef5cc7f8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx @@ -0,0 +1,38 @@ +import { useState } from "react"; +import { Button } from "antd"; +import { Eye, EyeOff } from "lucide-react"; + +export default function RedactableField({ + defaultHidden = true, + value, +}: { + defaultHidden?: boolean; + value: string | null; +}) { + const [isHidden, setIsHidden] = useState(defaultHidden); + + return ( +
+ + {value ? ( + isHidden ? ( + "•".repeat(value.length) + ) : ( + value + ) + ) : ( + Not configured + )} + + {value && ( +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index 0ece207662c..975a3bc7d78 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -3,11 +3,14 @@ import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Badge, Button, Card, Descriptions, Space, Typography } from "antd"; -import { Shield, Trash2 } from "lucide-react"; +import { Shield, Trash2, Edit } from "lucide-react"; import { useState } from "react"; import AddSSOSettingsModal from "./Modals/AddSSOSettingsModal"; import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; +import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import RedactableField from "./RedactableField"; +import { ssoProviderLogoMap, ssoProviderDisplayNames } from "./constants"; const { Title, Text } = Typography; @@ -16,6 +19,7 @@ export default function SSOSettings() { const { accessToken } = useAuthorized(); const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false); const [isAddModalVisible, setIsAddModalVisible] = useState(false); + const [isEditModalVisible, setIsEditModalVisible] = useState(false); const isSSOConfigured = Boolean(ssoSettings?.values.google_client_id) || Boolean(ssoSettings?.values.microsoft_client_id) || @@ -39,12 +43,6 @@ export default function SSOSettings() { } } - const renderRedactedValue = (value?: string | null) => ( - - {value ? "••••••••••••••••••••••••••••••••" : Not configured} - - ); - const renderEndpointValue = (value?: string | null) => ( {value || Not configured} @@ -67,44 +65,44 @@ export default function SSOSettings() { const providerConfigs = { google: { - providerText: "Google OAuth", + providerText: ssoProviderDisplayNames.google, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.google_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.google_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, ], }, microsoft: { - providerText: "Microsoft OAuth", + providerText: ssoProviderDisplayNames.microsoft, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.microsoft_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.microsoft_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Tenant", render: (values: any) => renderSimpleValue(values.microsoft_tenant) }, { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, ], }, okta: { - providerText: "Okta/Auth0", + providerText: ssoProviderDisplayNames.okta, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Authorization Endpoint", @@ -122,15 +120,15 @@ export default function SSOSettings() { ], }, generic: { - providerText: "Generic OAuth", + providerText: ssoProviderDisplayNames.generic, fields: [ { - label: "Client ID (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_id), + label: "Client ID", + render: (values: SSOSettingsValues) => , }, { - label: "Client Secret (Redacted)", - render: (values: SSOSettingsValues) => renderRedactedValue(values.generic_client_secret), + label: "Client Secret", + render: (values: SSOSettingsValues) => , }, { label: "Authorization Endpoint", @@ -160,7 +158,16 @@ export default function SSOSettings() { return ( - +
+ {ssoProviderLogoMap[selectedProvider] && ( + {selectedProvider} + )} + {config.providerText} +
{config.fields.map((field, index) => ( @@ -186,9 +193,14 @@ export default function SSOSettings() {
{isSSOConfigured && ( - + <> + + + )}
@@ -214,7 +226,15 @@ export default function SSOSettings() { setIsAddModalVisible(false); refetch(); }} - accessToken={accessToken} + /> + + setIsEditModalVisible(false)} + onSuccess={() => { + setIsEditModalVisible(false); + refetch(); + }} /> ); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts new file mode 100644 index 00000000000..595a961401a --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts @@ -0,0 +1,15 @@ +// SSO Provider logos +export const ssoProviderLogoMap: Record = { + google: "https://artificialanalysis.ai/img/logos/google_small.svg", + microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", + okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png", + generic: "", +}; + +// SSO Provider display names (consistent between select dropdown and table) +export const ssoProviderDisplayNames: Record = { + google: "Google SSO", + microsoft: "Microsoft SSO", + okta: "Okta / Auth0 SSO", + generic: "Generic SSO", +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts new file mode 100644 index 00000000000..1c878d7b7b3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts @@ -0,0 +1,274 @@ +import { processSSOSettingsPayload } from "./utils"; +import { describe, it, expect } from "vitest"; + +describe("processSSOSettingsPayload", () => { + describe("without role mappings", () => { + it("should return all fields except role mapping fields when use_role_mappings is false", () => { + const formValues = { + proxy_admin_teams: "team1, team2", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: false, + other_field: "value", + another_field: 123, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + another_field: 123, + }); + expect(result.role_mappings).toBeUndefined(); + }); + + it("should return all fields except role mapping fields when use_role_mappings is not present", () => { + const formValues = { + proxy_admin_teams: "team1", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + }); + expect(result.role_mappings).toBeUndefined(); + }); + }); + + describe("with role mappings enabled", () => { + it("should create role mappings with all team types populated", () => { + const formValues = { + proxy_admin_teams: "admin1, admin2", + admin_viewer_teams: "viewer1, viewer2, viewer3", + internal_user_teams: "user1", + internal_viewer_teams: "internal_viewer1, internal_viewer2", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.other_field).toBe("value"); + expect(result.role_mappings).toEqual({ + provider: "generic", + group_claim: "groups", + default_role: "proxy_admin", + roles: { + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1", "viewer2", "viewer3"], + internal_user: ["user1"], + internal_user_viewer: ["internal_viewer1", "internal_viewer2"], + }, + }); + }); + + it("should handle empty team strings", () => { + const formValues = { + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle undefined team fields", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle whitespace-only team strings", () => { + const formValues = { + proxy_admin_teams: " ", + admin_viewer_teams: ", , ,", + internal_user_teams: "user1, , user2", + internal_viewer_teams: "viewer1, ,viewer2", + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should trim whitespace from team names", () => { + const formValues = { + proxy_admin_teams: " admin1 , admin2 ", + admin_viewer_teams: " viewer1 ", + internal_user_teams: " user1 , user2 ", + internal_viewer_teams: "viewer1,viewer2", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1"], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should filter out empty strings after trimming", () => { + const formValues = { + proxy_admin_teams: "admin1,,admin2, , admin3", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles.proxy_admin).toEqual(["admin1", "admin2", "admin3"]); + }); + }); + + describe("default role mapping", () => { + it("should map internal_user_viewer correctly", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user_viewer"); + }); + + it("should map internal_user correctly", () => { + const formValues = { + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should map proxy_admin_viewer correctly", () => { + const formValues = { + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin_viewer"); + }); + + it("should map proxy_admin correctly", () => { + const formValues = { + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin"); + }); + + it("should default to internal_user for unknown roles", () => { + const formValues = { + default_role: "unknown_role", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should default to internal_user for undefined default_role", () => { + const formValues = { + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + }); + + describe("edge cases", () => { + it("should handle empty form values", () => { + const result = processSSOSettingsPayload({}); + + expect(result).toEqual({}); + }); + + it("should preserve other fields in the payload", () => { + const formValues = { + use_role_mappings: false, + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts new file mode 100644 index 00000000000..3533e1226c2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts @@ -0,0 +1,54 @@ +/** + * Processes SSO settings form values and transforms them into the payload format expected by the API + * Handles role mappings transformation and field extraction + */ +export const processSSOSettingsPayload = (formValues: Record): Record => { + const { + proxy_admin_teams, + admin_viewer_teams, + internal_user_teams, + internal_viewer_teams, + default_role, + group_claim, + use_role_mappings, + ...rest + } = formValues; + + const payload: any = { + ...rest, + }; + + // Add role mappings if use_role_mappings is checked + if (use_role_mappings) { + // Helper function to split comma-separated string into array + const splitTeams = (teams: string | undefined): string[] => { + if (!teams || teams.trim() === "") return []; + return teams + .split(",") + .map((team) => team.trim()) + .filter((team) => team.length > 0); + }; + + // Map default role display values to backend values + const defaultRoleMapping: Record = { + internal_user_viewer: "internal_user_viewer", + internal_user: "internal_user", + proxy_admin_viewer: "proxy_admin_viewer", + proxy_admin: "proxy_admin", + }; + + payload.role_mappings = { + provider: "generic", + group_claim, + default_role: defaultRoleMapping[default_role] || "internal_user", + roles: { + proxy_admin: splitTeams(proxy_admin_teams), + proxy_admin_viewer: splitTeams(admin_viewer_teams), + internal_user: splitTeams(internal_user_teams), + internal_user_viewer: splitTeams(internal_viewer_teams), + }, + }; + } + + return payload; +}; From 5c7523b11e3bf1786addd39ed71246f2c8c08b50 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 12:43:13 -0800 Subject: [PATCH 243/330] Fixing tests --- .../Modals/AddSSOSettingsModal.test.tsx | 23 +++- .../SSOSettings/RedactableField.test.tsx | 108 ++++++++++++++++++ .../AdminSettings/SSOSettings/SSOSettings.tsx | 8 +- 3 files changed, 133 insertions(+), 6 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx index aae28191031..e5a5af6cba6 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.test.tsx @@ -1,5 +1,6 @@ -import { render, screen } from "@testing-library/react"; +import { screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../../../tests/test-utils"; import AddSSOSettingsModal from "./AddSSOSettingsModal"; // Mock networking functions @@ -12,12 +13,30 @@ vi.mock("@/components/shared/errorUtils", () => ({ parseErrorMessage: vi.fn((error) => error?.message || "Unknown error"), })); +// Mock the useAuthorized hook +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: "test-access-token", + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "admin", + }), +})); + +// Mock NotificationsManager +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + describe("AddSSOSettingsModal", () => { it("should render", () => { const onCancel = vi.fn(); const onSuccess = vi.fn(); - render(); + renderWithProviders(); expect(screen.getByText("SSO Provider")).toBeInTheDocument(); expect(screen.getByText("Cancel")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx new file mode 100644 index 00000000000..a047d7aea4f --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx @@ -0,0 +1,108 @@ +import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import RedactableField from "./RedactableField"; + +describe("RedactableField", () => { + describe("when value is null", () => { + it("should display 'Not configured' text", () => { + render(); + + expect(screen.getByText("Not configured")).toBeInTheDocument(); + }); + + it("should not display toggle button", () => { + render(); + + // There should be no button elements + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + }); + + describe("when value is provided", () => { + const testValue = "secret-password"; + + it("should be hidden by default and show redacted dots", () => { + render(); + + // Should show dots equal to the length of the value + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should show actual value when defaultHidden is false", () => { + render(); + + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + }); + + it("should display toggle button with eye icon when hidden", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // Check that the Eye icon is rendered (we can check by title or by the presence of the icon) + // The button should contain the Eye icon when hidden + const eyeIcon = button.querySelector("svg"); + expect(eyeIcon).toBeInTheDocument(); + }); + + it("should display toggle button with eye-off icon when shown", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // The button should contain the EyeOff icon when shown + const eyeOffIcon = button.querySelector("svg"); + expect(eyeOffIcon).toBeInTheDocument(); + }); + + it("should toggle visibility when button is clicked", () => { + render(); + + // Initially hidden + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + + // Click to show + const button = screen.getByRole("button"); + fireEvent.click(button); + + // Should now show the actual value + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + + // Click again to hide + fireEvent.click(button); + + // Should be hidden again + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should handle empty string value", () => { + render(); + + // Empty string should show "Not configured" since value is falsy + expect(screen.getByText("Not configured")).toBeInTheDocument(); + + // No toggle button for empty string + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + + it("should handle different value lengths correctly", () => { + const shortValue = "hi"; + const longValue = "this-is-a-very-long-secret-value"; + + const { rerender } = render(); + expect(screen.getByText("••")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("•".repeat(longValue.length))).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index 975a3bc7d78..d339f6d0e36 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -2,15 +2,15 @@ import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { Badge, Button, Card, Descriptions, Space, Typography } from "antd"; -import { Shield, Trash2, Edit } from "lucide-react"; +import { Button, Card, Descriptions, Space, Typography } from "antd"; +import { Edit, Shield, Trash2 } from "lucide-react"; import { useState } from "react"; import AddSSOSettingsModal from "./Modals/AddSSOSettingsModal"; import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; -import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; import RedactableField from "./RedactableField"; -import { ssoProviderLogoMap, ssoProviderDisplayNames } from "./constants"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants"; const { Title, Text } = Typography; From 2cbcaf2abf463a137e9a7f164c305da15e6c9413 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 13:34:17 -0800 Subject: [PATCH 244/330] refactor(utils): lazy load heavy imports to improve import time and memory usage (#18610) * refactor(utils): lazy load heavy imports to improve import time - Move BaseVectorStore, CredentialAccessor, and exception_mapping_utils imports to lazy loading via __getattr__ - Add _get_utils_globals() helper function following pattern from _lazy_imports.py - Refactor __getattr__ to use consistent caching pattern matching __init__.py - Update load_credentials_from_list to use lazy-loaded CredentialAccessor This reduces import time and memory usage by only loading these modules when they're actually accessed, not during module import. * refactor(utils): lazy load additional heavy imports to improve import time - Move get_llm_provider, _is_non_openai_azure_model to lazy loading - Move get_supported_openai_params to lazy loading - Move convert_dict_to_response functions (LiteLLMResponseObjectHandler, convert_to_model_response_object, etc.) to lazy loading - Move get_api_base and ResponseMetadata to lazy loading - Move _parse_content_for_reasoning to lazy loading - Update all internal usages to access via getattr(sys.modules[__name__], ...) This reduces import time and memory usage by only loading these modules when they're actually accessed, not during module import. * fix(utils): suppress PLR0915 linter warning for __getattr__ function The __getattr__ function intentionally has many statements to handle multiple lazy-loaded imports. Add noqa comment to suppress the warning. * fix(utils): add type stubs for lazy-loaded functions in TYPE_CHECKING block Add type imports and declarations in TYPE_CHECKING block to help mypy understand the types of lazy-loaded functions accessed via __getattr__. This follows the same pattern used in __init__.py for lazy-loaded items. --- litellm/utils.py | 249 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 209 insertions(+), 40 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 3dbeeb970a4..e5eb57b0712 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -71,9 +71,6 @@ from litellm.constants import ( OPENAI_EMBEDDING_PARAMS, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.integrations.vector_store_integrations.base_vector_store import ( - BaseVectorStore, -) # Import cached imports utilities from litellm.litellm_core_utils.cached_imports import ( @@ -86,48 +83,21 @@ from litellm.litellm_core_utils.core_helpers import ( map_finish_reason, process_response_headers, ) -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dot_notation_indexing import ( delete_nested_value, is_nested_path, ) -from litellm.litellm_core_utils.exception_mapping_utils import ( - _get_response_headers, - exception_type, - get_error_message, -) from litellm.litellm_core_utils.get_litellm_params import ( _get_base_model_from_litellm_call_metadata, get_litellm_params, ) -from litellm.litellm_core_utils.get_llm_provider_logic import ( - _is_non_openai_azure_model, - get_llm_provider, -) -from litellm.litellm_core_utils.get_supported_openai_params import ( - get_supported_openai_params, -) from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe -from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( - LiteLLMResponseObjectHandler, - _handle_invalid_parallel_tool_calls, - convert_to_model_response_object, - convert_to_streaming_response, - convert_to_streaming_response_async, -) -from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( get_formatted_prompt, ) from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - ResponseMetadata, -) -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - _parse_content_for_reasoning, -) from litellm.litellm_core_utils.redact_messages import ( LiteLLMLoggingObject, redact_message_input_output_from_logging, @@ -346,6 +316,30 @@ if TYPE_CHECKING: from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.proxy._types import AllowedModelRegion + # Type stubs for lazy-loaded functions to help mypy understand their types + # These imports allow mypy to understand the types when these are accessed via __getattr__ + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model, + get_llm_provider, + ) + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, + ) + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler, + _handle_invalid_parallel_tool_calls, + convert_to_model_response_object, + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning, + ) from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -618,10 +612,22 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]: return applied_guardrails +def _get_utils_globals() -> dict: + """ + Get the globals dictionary of the utils module. + + This is where we cache imported attributes so we don't import them twice. + """ + return sys.modules[__name__].__dict__ + + def load_credentials_from_list(kwargs: dict): """ Updates kwargs with the credentials if credential_name in kwarg """ + # Access CredentialAccessor via module to trigger lazy loading if needed + CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -2259,6 +2265,7 @@ def supports_response_schema( """ ## GET LLM PROVIDER ## try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) @@ -2956,6 +2963,9 @@ def get_optional_params_embeddings( # noqa: PLR0915 additional_drop_params: Optional[List[str]] = None, **kwargs, ): + # Lazy load get_supported_openai_params + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -3758,6 +3768,7 @@ def get_optional_params( # noqa: PLR0915 message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -4895,6 +4906,7 @@ def get_max_tokens(model: str) -> Optional[int]: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) @@ -5015,6 +5027,7 @@ def _get_potential_model_names( if custom_llm_provider is None: # Get custom_llm_provider try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5737,6 +5750,7 @@ def validate_environment( # noqa: PLR0915 } ## EXTRACT LLM PROVIDER - if model name provided try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6299,6 +6313,7 @@ def register_prompt_template( complete_model = model potential_models = [complete_model] try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -6384,6 +6399,7 @@ class TextCompletionStreamWrapper: except StopIteration: raise StopIteration except Exception as e: + exception_type = getattr(sys.modules[__name__], 'exception_type') raise exception_type( model=self.model, custom_llm_provider=self.custom_llm_provider or "", @@ -8705,16 +8721,169 @@ def should_run_mock_completion( return False -# Re-export encoding from main.py for backward compatibility -# This allows tests to import: from litellm.utils import encoding -# We use a lazy import to avoid loading main.py at utils.py import time -def __getattr__(name: str) -> Any: +def __getattr__(name: str) -> Any: # noqa: PLR0915 """Lazy import handler for utils module""" + _globals = _get_utils_globals() + + # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": - # Cache it in the module's __dict__ for subsequent accesses - import sys - - from litellm.main import encoding as _encoding - sys.modules[__name__].__dict__["encoding"] = _encoding - return _encoding + # Check if already cached + if "encoding" not in _globals: + from litellm.main import encoding as _encoding + _globals["encoding"] = _encoding + return _globals["encoding"] + + # Lazy load BaseVectorStore to avoid loading it at module import time + if name == "BaseVectorStore": + # Check if already cached + if "BaseVectorStore" not in _globals: + from litellm.integrations.vector_store_integrations.base_vector_store import ( + BaseVectorStore as _BaseVectorStore, + ) + _globals["BaseVectorStore"] = _BaseVectorStore + return _globals["BaseVectorStore"] + + # Lazy load CredentialAccessor to avoid loading it at module import time + if name == "CredentialAccessor": + # Check if already cached + if "CredentialAccessor" not in _globals: + from litellm.litellm_core_utils.credential_accessor import ( + CredentialAccessor as _CredentialAccessor, + ) + _globals["CredentialAccessor"] = _CredentialAccessor + return _globals["CredentialAccessor"] + + # Lazy load exception_mapping_utils functions to avoid loading at module import time + if name == "exception_type": + # Check if already cached + if "exception_type" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + exception_type as _exception_type, + ) + _globals["exception_type"] = _exception_type + return _globals["exception_type"] + + if name == "get_error_message": + # Check if already cached + if "get_error_message" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + get_error_message as _get_error_message, + ) + _globals["get_error_message"] = _get_error_message + return _globals["get_error_message"] + + if name == "_get_response_headers": + # Check if already cached + if "_get_response_headers" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers as __get_response_headers, + ) + _globals["_get_response_headers"] = __get_response_headers + return _globals["_get_response_headers"] + + # Lazy load get_llm_provider_logic functions to avoid loading at module import time + if name == "get_llm_provider": + # Check if already cached + if "get_llm_provider" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider as _get_llm_provider, + ) + _globals["get_llm_provider"] = _get_llm_provider + return _globals["get_llm_provider"] + + if name == "_is_non_openai_azure_model": + # Check if already cached + if "_is_non_openai_azure_model" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model as __is_non_openai_azure_model, + ) + _globals["_is_non_openai_azure_model"] = __is_non_openai_azure_model + return _globals["_is_non_openai_azure_model"] + + # Lazy load get_supported_openai_params to avoid loading at module import time + if name == "get_supported_openai_params": + # Check if already cached + if "get_supported_openai_params" not in _globals: + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params as _get_supported_openai_params, + ) + _globals["get_supported_openai_params"] = _get_supported_openai_params + return _globals["get_supported_openai_params"] + + # Lazy load convert_dict_to_response functions to avoid loading at module import time + if name == "LiteLLMResponseObjectHandler": + # Check if already cached + if "LiteLLMResponseObjectHandler" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler as _LiteLLMResponseObjectHandler, + ) + _globals["LiteLLMResponseObjectHandler"] = _LiteLLMResponseObjectHandler + return _globals["LiteLLMResponseObjectHandler"] + + if name == "_handle_invalid_parallel_tool_calls": + # Check if already cached + if "_handle_invalid_parallel_tool_calls" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _handle_invalid_parallel_tool_calls as __handle_invalid_parallel_tool_calls, + ) + _globals["_handle_invalid_parallel_tool_calls"] = __handle_invalid_parallel_tool_calls + return _globals["_handle_invalid_parallel_tool_calls"] + + if name == "convert_to_model_response_object": + # Check if already cached + if "convert_to_model_response_object" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object as _convert_to_model_response_object, + ) + _globals["convert_to_model_response_object"] = _convert_to_model_response_object + return _globals["convert_to_model_response_object"] + + if name == "convert_to_streaming_response": + # Check if already cached + if "convert_to_streaming_response" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response as _convert_to_streaming_response, + ) + _globals["convert_to_streaming_response"] = _convert_to_streaming_response + return _globals["convert_to_streaming_response"] + + if name == "convert_to_streaming_response_async": + # Check if already cached + if "convert_to_streaming_response_async" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async as _convert_to_streaming_response_async, + ) + _globals["convert_to_streaming_response_async"] = _convert_to_streaming_response_async + return _globals["convert_to_streaming_response_async"] + + # Lazy load get_api_base to avoid loading at module import time + if name == "get_api_base": + # Check if already cached + if "get_api_base" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_api_base import ( + get_api_base as _get_api_base, + ) + _globals["get_api_base"] = _get_api_base + return _globals["get_api_base"] + + # Lazy load ResponseMetadata to avoid loading at module import time + if name == "ResponseMetadata": + # Check if already cached + if "ResponseMetadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata as _ResponseMetadata, + ) + _globals["ResponseMetadata"] = _ResponseMetadata + return _globals["ResponseMetadata"] + + # Lazy load _parse_content_for_reasoning to avoid loading at module import time + if name == "_parse_content_for_reasoning": + # Check if already cached + if "_parse_content_for_reasoning" not in _globals: + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning as __parse_content_for_reasoning, + ) + _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning + return _globals["_parse_content_for_reasoning"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From 1c5c303e986bb5c11756ab411713badd6c1a1362 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 14:07:57 -0800 Subject: [PATCH 245/330] refactor(utils): implement lazy loading for provider configs, model info classes, streaming handlers, and redact utilities (#18611) * refactor(utils): lazy load redact_messages imports to improve import time - Move LiteLLMLoggingObject and redact_message_input_output_from_logging to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - These are only used in type annotations (with from __future__ import annotations), so lazy loading works correctly This reduces import time by deferring the redact_messages module import until these are actually accessed. * refactor(utils): lazy load CustomStreamWrapper to improve import time - Move CustomStreamWrapper from streaming_handler to lazy loading via __getattr__ - Add type stub in TYPE_CHECKING block for mypy type checking - CustomStreamWrapper is not used internally in utils.py, only exported for other modules This reduces import time by deferring the streaming_handler module import until CustomStreamWrapper is actually accessed. * refactor(utils): lazy load BaseGoogleGenAIGenerateContentConfig to improve import time - Move BaseGoogleGenAIGenerateContentConfig from google_genai.transformation to lazy loading via __getattr__ - Add type stub in TYPE_CHECKING block for mypy type checking - BaseGoogleGenAIGenerateContentConfig is only used in type annotations (with from __future__ import annotations), so lazy loading works correctly This reduces import time by deferring the google_genai.transformation module import until BaseGoogleGenAIGenerateContentConfig is actually accessed. * refactor(utils): lazy load BaseOCRConfig, BaseSearchConfig, and BaseTextToSpeechConfig - Move BaseOCRConfig, BaseSearchConfig, and BaseTextToSpeechConfig to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - These config classes are only used in quoted type annotations (forward references), so lazy loading works correctly This reduces import time by deferring the transformation module imports until these config classes are actually accessed. * refactor(utils): lazy load BedrockModelInfo, CohereModelInfo, and MistralOCRConfig - Move BedrockModelInfo, CohereModelInfo, and MistralOCRConfig to lazy loading via __getattr__ - Add type stubs in TYPE_CHECKING block for mypy type checking - Update internal usages to use getattr pattern for accessing lazy-loaded classes - These provider-specific model info classes are only used in specific code paths, so lazy loading reduces initial import time This reduces import time by deferring the bedrock, cohere, and mistral module imports until these classes are actually accessed. * fix(utils): remove duplicate MistralOCRConfig import in TYPE_CHECKING block --- litellm/utils.py | 130 ++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 116 insertions(+), 14 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index e5eb57b0712..d373b5102ef 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -98,22 +98,8 @@ from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) -from litellm.litellm_core_utils.redact_messages import ( - LiteLLMLoggingObject, - redact_message_input_output_from_logging, -) from litellm.litellm_core_utils.rules import Rules -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.google_genai.transformation import ( - BaseGoogleGenAIGenerateContentConfig, -) -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig -from litellm.llms.base_llm.search.transformation import BaseSearchConfig -from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig -from litellm.llms.bedrock.common_utils import BedrockModelInfo -from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, reset_retry_policy, @@ -340,6 +326,20 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.prompt_templates.common_utils import ( _parse_content_for_reasoning, ) + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject, + redact_message_input_output_from_logging, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, + ) + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig + from litellm.llms.base_llm.search.transformation import BaseSearchConfig + from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig + from litellm.llms.bedrock.common_utils import BedrockModelInfo + from litellm.llms.cohere.common_utils import CohereModelInfo + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -4023,6 +4023,7 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "bedrock": + BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo') bedrock_route = BedrockModelInfo.get_bedrock_route(model) bedrock_base_model = BedrockModelInfo.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": @@ -7440,6 +7441,7 @@ class ProviderConfigManager: litellm.LlmProviders.COHERE_CHAT == provider or litellm.LlmProviders.COHERE == provider ): + CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') route = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -8290,6 +8292,7 @@ class ProviderConfigManager: return get_vertex_ai_ocr_config(model=model) + MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8886,4 +8889,103 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning return _globals["_parse_content_for_reasoning"] + # Lazy load redact_messages to avoid loading at module import time + if name == "LiteLLMLoggingObject": + # Check if already cached + if "LiteLLMLoggingObject" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject as _LiteLLMLoggingObject, + ) + _globals["LiteLLMLoggingObject"] = _LiteLLMLoggingObject + return _globals["LiteLLMLoggingObject"] + + if name == "redact_message_input_output_from_logging": + # Check if already cached + if "redact_message_input_output_from_logging" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging as _redact_message_input_output_from_logging, + ) + _globals["redact_message_input_output_from_logging"] = _redact_message_input_output_from_logging + return _globals["redact_message_input_output_from_logging"] + + # Lazy load CustomStreamWrapper to avoid loading at module import time + if name == "CustomStreamWrapper": + # Check if already cached + if "CustomStreamWrapper" not in _globals: + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper as _CustomStreamWrapper, + ) + _globals["CustomStreamWrapper"] = _CustomStreamWrapper + return _globals["CustomStreamWrapper"] + + # Lazy load BaseGoogleGenAIGenerateContentConfig to avoid loading at module import time + if name == "BaseGoogleGenAIGenerateContentConfig": + # Check if already cached + if "BaseGoogleGenAIGenerateContentConfig" not in _globals: + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig as _BaseGoogleGenAIGenerateContentConfig, + ) + _globals["BaseGoogleGenAIGenerateContentConfig"] = _BaseGoogleGenAIGenerateContentConfig + return _globals["BaseGoogleGenAIGenerateContentConfig"] + + # Lazy load BaseOCRConfig to avoid loading at module import time + if name == "BaseOCRConfig": + # Check if already cached + if "BaseOCRConfig" not in _globals: + from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig as _BaseOCRConfig, + ) + _globals["BaseOCRConfig"] = _BaseOCRConfig + return _globals["BaseOCRConfig"] + + # Lazy load BaseSearchConfig to avoid loading at module import time + if name == "BaseSearchConfig": + # Check if already cached + if "BaseSearchConfig" not in _globals: + from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig as _BaseSearchConfig, + ) + _globals["BaseSearchConfig"] = _BaseSearchConfig + return _globals["BaseSearchConfig"] + + # Lazy load BaseTextToSpeechConfig to avoid loading at module import time + if name == "BaseTextToSpeechConfig": + # Check if already cached + if "BaseTextToSpeechConfig" not in _globals: + from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig as _BaseTextToSpeechConfig, + ) + _globals["BaseTextToSpeechConfig"] = _BaseTextToSpeechConfig + return _globals["BaseTextToSpeechConfig"] + + # Lazy load BedrockModelInfo to avoid loading at module import time + if name == "BedrockModelInfo": + # Check if already cached + if "BedrockModelInfo" not in _globals: + from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo as _BedrockModelInfo, + ) + _globals["BedrockModelInfo"] = _BedrockModelInfo + return _globals["BedrockModelInfo"] + + # Lazy load CohereModelInfo to avoid loading at module import time + if name == "CohereModelInfo": + # Check if already cached + if "CohereModelInfo" not in _globals: + from litellm.llms.cohere.common_utils import ( + CohereModelInfo as _CohereModelInfo, + ) + _globals["CohereModelInfo"] = _CohereModelInfo + return _globals["CohereModelInfo"] + + # Lazy load MistralOCRConfig to avoid loading at module import time + if name == "MistralOCRConfig": + # Check if already cached + if "MistralOCRConfig" not in _globals: + from litellm.llms.mistral.ocr.transformation import ( + MistralOCRConfig as _MistralOCRConfig, + ) + _globals["MistralOCRConfig"] = _MistralOCRConfig + return _globals["MistralOCRConfig"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From a6c3fb1fb598422e251fa5bd13e2aa4d16c7d3a6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 15:10:52 -0800 Subject: [PATCH 246/330] E2E test see models for specific provider --- ui/litellm-dashboard/e2e_tests/constants.ts | 1 + .../tests/modelsPage/addModel.spec.ts | 23 +++++++++++++++++++ .../tests/navigation/sidebar.spec.ts | 3 ++- 3 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/e2e_tests/constants.ts create mode 100644 ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts new file mode 100644 index 00000000000..b07bd68fcf1 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/constants.ts @@ -0,0 +1 @@ +export const ADMIN_STORAGE_PATH = "admin.storageState.json"; diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts new file mode 100644 index 00000000000..5fa11a98ef6 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -0,0 +1,23 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; + +test.describe("Add Model", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to see all models for a specific provider in the model dropdown", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + + await page.getByText("Models + Endpoints").click(); + await page.getByRole("tab", { name: "Add Model" }).click(); + + const providerInputDropdown = page.getByRole("combobox", { name: /Provider/i }); + await providerInputDropdown.fill("Anthropic"); + await page.waitForTimeout(1000); + await providerInputDropdown.press("Enter"); + await page.waitForTimeout(1000); + + const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first(); + await providerModelsDropdown.click(); + await expect(page.getByTitle("claude-haiku-4-5", { exact: true })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index dafb03a7cbd..6801f891e87 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -1,5 +1,6 @@ import test, { expect } from "@playwright/test"; import { Role } from "../../fixtures/roles"; +import { ADMIN_STORAGE_PATH } from "../../constants"; const sidebarButtons = { [Role.ProxyAdmin]: [ @@ -16,7 +17,7 @@ const sidebarButtons = { ], }; -const roles = [{ role: Role.ProxyAdmin, storage: "admin.storageState.json" }]; +const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; for (const { role, storage } of roles) { test.describe(`${role} sidebar`, () => { From dd1ccec7348b73f9a3d7d1ec8ac5fd29651de759 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 15:50:53 -0800 Subject: [PATCH 247/330] refactor(utils): lazy load 15 additional imports to improve import time (#18613) * refactor(utils): lazy load 15 additional imports to improve import time - Move Rules, AsyncHTTPHandler, HTTPHandler to lazy loading via __getattr__ - Move get_num_retries_from_retry_policy, reset_retry_policy to lazy loading - Move get_secret to lazy loading - Move cached_imports functions (get_coroutine_checker, get_litellm_logging_class, get_set_callbacks) to lazy loading - Move core_helpers functions (get_litellm_metadata_from_kwargs, map_finish_reason, process_response_headers) to lazy loading - Move dot_notation_indexing functions (delete_nested_value, is_nested_path) to lazy loading - Move get_litellm_params functions to lazy loading - Move _ensure_extra_body_is_safe, get_formatted_prompt, get_response_headers, update_response_metadata to lazy loading - Move executor to lazy loading - Move BaseAnthropicMessagesConfig, BaseAudioTranscriptionConfig to lazy loading - Add type stubs in TYPE_CHECKING block for mypy type checking - These functions/classes are exported for other modules but not used internally in utils.py, so lazy loading is safe and improves startup performance * fix(utils): use getattr for Rules and get_coroutine_checker in client decorator - Update Rules() instantiation in client decorator to use getattr for lazy loading - Update Rules.has_pre_call_rules() usage in function_setup to use getattr - Update get_coroutine_checker() usage in client decorator to use getattr - Fixes NameError: name 'Rules' is not defined error that occurs when Rules is lazy-loaded * fix(utils): use getattr for get_litellm_logging_class in function_setup - Update get_litellm_logging_class() usage in function_setup to use getattr for lazy loading - Fixes NameError: name 'get_litellm_logging_class' is not defined error that occurs when get_litellm_logging_class is lazy-loaded * fix(utils): use getattr for get_set_callbacks in function_setup - Update get_set_callbacks() usage in function_setup to use getattr for lazy loading - Fixes NameError: name 'get_set_callbacks' is not defined error that occurs when get_set_callbacks is lazy-loaded * fix(utils): use getattr for all lazy-loaded imports in utils.py - Update update_response_metadata (4 occurrences) to use getattr - Update executor.submit (1 occurrence) to use getattr - Update get_num_retries_from_retry_policy (2 occurrences) to use getattr - Update reset_retry_policy (2 occurrences) to use getattr - Update is_nested_path and delete_nested_value (1 occurrence each) to use getattr - Update _ensure_extra_body_is_safe (1 occurrence) to use getattr Fixes NameError errors that occur when these functions/classes are lazy-loaded but used directly in utils.py * fix(utils): use getattr for _get_base_model_from_litellm_call_metadata in _get_base_model_from_metadata - Update _get_base_model_from_litellm_call_metadata usage to use getattr for lazy loading - Fixes NameError: name '_get_base_model_from_litellm_call_metadata' is not defined * fix(utils): use getattr for second _get_base_model_from_litellm_call_metadata usage - Fix the second occurrence of _get_base_model_from_litellm_call_metadata on line 7052 - Both occurrences in _get_base_model_from_metadata now use getattr for lazy loading * fix(utils): fix indentation in _get_base_model_from_metadata function * fix(utils): use getattr for get_litellm_metadata_from_kwargs in _get_litellm_params - Update get_litellm_metadata_from_kwargs usage to use getattr for lazy loading - Fixes NameError: name 'get_litellm_metadata_from_kwargs' is not defined * fix(utils): fix syntax error in get_litellm_metadata_from_kwargs fix - Move getattr call before cast statement to fix syntax error --- litellm/utils.py | 330 ++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 284 insertions(+), 46 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d373b5102ef..6f4652c8278 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -72,39 +72,7 @@ from litellm.constants import ( TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -# Import cached imports utilities -from litellm.litellm_core_utils.cached_imports import ( - get_coroutine_checker, - get_litellm_logging_class, - get_set_callbacks, -) -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - map_finish_reason, - process_response_headers, -) -from litellm.litellm_core_utils.dot_notation_indexing import ( - delete_nested_value, - is_nested_path, -) -from litellm.litellm_core_utils.get_litellm_params import ( - _get_base_model_from_litellm_call_metadata, - get_litellm_params, -) -from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe -from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( - get_formatted_prompt, -) -from litellm.litellm_core_utils.llm_response_utils.get_headers import ( - get_response_headers, -) -from litellm.litellm_core_utils.rules import Rules -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.router_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy, - reset_retry_policy, -) -from litellm.secret_managers.main import get_secret + _CachingHandlerResponse = None _LLMCachingHandler = None @@ -280,16 +248,7 @@ from typing import ( from openai import OpenAIError as OriginalError -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - update_response_metadata, -) -from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.anthropic_messages.transformation import ( - BaseAnthropicMessagesConfig, -) -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) +# These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, type_to_response_format_param, @@ -340,6 +299,49 @@ if TYPE_CHECKING: from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + # Type stubs for lazy-loaded functions and classes + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker, + get_litellm_logging_class, + get_set_callbacks, + ) + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + map_finish_reason, + process_response_headers, + ) + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value, + is_nested_path, + ) + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata, + get_litellm_params, + ) + from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt, + ) + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, + ) + from litellm.litellm_core_utils.rules import Rules + from litellm.litellm_core_utils.thread_pool_executor import executor + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig, + ) + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy, + reset_retry_policy, + ) + from litellm.secret_managers.main import get_secret from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -816,6 +818,7 @@ def function_setup( # noqa: PLR0915 + litellm.failure_callback ) ) + get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks') get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS if len(litellm.input_callback) > 0: @@ -943,6 +946,7 @@ def function_setup( # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### + Rules = getattr(sys.modules[__name__], 'Rules') if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -1075,6 +1079,7 @@ def function_setup( # noqa: PLR0915 call_type=call_type, ): stream = True + get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class') logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1158,6 +1163,8 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), @@ -1343,6 +1350,7 @@ def post_call_processing( def client(original_function): # noqa: PLR0915 + Rules = getattr(sys.modules[__name__], 'Rules') rules_obj = Rules() @wraps(original_function) @@ -1528,6 +1536,7 @@ def client(original_function): # noqa: PLR0915 ) else: # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1571,6 +1580,7 @@ def client(original_function): # noqa: PLR0915 # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx = contextvars.copy_context() + executor = getattr(sys.modules[__name__], 'executor') executor.submit( ctx.run, logging_obj.success_handler, @@ -1579,6 +1589,7 @@ def client(original_function): # noqa: PLR0915 end_time, ) # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1595,6 +1606,8 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), @@ -1766,6 +1779,7 @@ def client(original_function): # noqa: PLR0915 chunks, messages=kwargs.get("messages", None) ) else: + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1830,6 +1844,7 @@ def client(original_function): # noqa: PLR0915 end_time=end_time, ) + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1905,6 +1920,7 @@ def client(original_function): # noqa: PLR0915 setattr(e, "timeout", timeout) raise e + get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') is_coroutine = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -4448,6 +4464,8 @@ def get_optional_params( # noqa: PLR0915 # Apply nested drops from additional_drop_params if additional_drop_params: + is_nested_path = getattr(sys.modules[__name__], 'is_nested_path') + delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value') nested_paths = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -4497,6 +4515,7 @@ def add_provider_specific_params_to_optional_params( else: processed_extra_body = initial_extra_body + _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe') optional_params["extra_body"] = _ensure_extra_body_is_safe( extra_body=processed_extra_body ) @@ -7022,14 +7041,14 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata = litellm_params.get("metadata", {}) - base_model_from_metadata = _get_base_model_from_litellm_call_metadata( - metadata=metadata - ) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') + base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) if base_model_from_metadata is not None: return base_model_from_metadata # Also check litellm_metadata (used by Responses API and other generic API calls) litellm_metadata = litellm_params.get("litellm_metadata", {}) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata) return None @@ -8433,6 +8452,7 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ + get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs') _metadata = cast( dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) ) @@ -8988,4 +9008,222 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["MistralOCRConfig"] = _MistralOCRConfig return _globals["MistralOCRConfig"] + # Lazy load Rules to avoid loading at module import time + if name == "Rules": + # Check if already cached + if "Rules" not in _globals: + from litellm.litellm_core_utils.rules import Rules as _Rules + _globals["Rules"] = _Rules + return _globals["Rules"] + + # Lazy load AsyncHTTPHandler and HTTPHandler to avoid loading at module import time + if name == "AsyncHTTPHandler": + # Check if already cached + if "AsyncHTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler as _AsyncHTTPHandler, + ) + _globals["AsyncHTTPHandler"] = _AsyncHTTPHandler + return _globals["AsyncHTTPHandler"] + + if name == "HTTPHandler": + # Check if already cached + if "HTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler as _HTTPHandler, + ) + _globals["HTTPHandler"] = _HTTPHandler + return _globals["HTTPHandler"] + + # Lazy load get_num_retries_from_retry_policy and reset_retry_policy to avoid loading at module import time + if name == "get_num_retries_from_retry_policy": + # Check if already cached + if "get_num_retries_from_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, + ) + _globals["get_num_retries_from_retry_policy"] = _get_num_retries_from_retry_policy + return _globals["get_num_retries_from_retry_policy"] + + if name == "reset_retry_policy": + # Check if already cached + if "reset_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + reset_retry_policy as _reset_retry_policy, + ) + _globals["reset_retry_policy"] = _reset_retry_policy + return _globals["reset_retry_policy"] + + # Lazy load get_secret to avoid loading at module import time + if name == "get_secret": + # Check if already cached + if "get_secret" not in _globals: + from litellm.secret_managers.main import get_secret as _get_secret + _globals["get_secret"] = _get_secret + return _globals["get_secret"] + + # Lazy load cached_imports functions to avoid loading at module import time + if name == "get_coroutine_checker": + # Check if already cached + if "get_coroutine_checker" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker as _get_coroutine_checker, + ) + _globals["get_coroutine_checker"] = _get_coroutine_checker + return _globals["get_coroutine_checker"] + + if name == "get_litellm_logging_class": + # Check if already cached + if "get_litellm_logging_class" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_litellm_logging_class as _get_litellm_logging_class, + ) + _globals["get_litellm_logging_class"] = _get_litellm_logging_class + return _globals["get_litellm_logging_class"] + + if name == "get_set_callbacks": + # Check if already cached + if "get_set_callbacks" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_set_callbacks as _get_set_callbacks, + ) + _globals["get_set_callbacks"] = _get_set_callbacks + return _globals["get_set_callbacks"] + + # Lazy load core_helpers functions to avoid loading at module import time + if name == "get_litellm_metadata_from_kwargs": + # Check if already cached + if "get_litellm_metadata_from_kwargs" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs as _get_litellm_metadata_from_kwargs, + ) + _globals["get_litellm_metadata_from_kwargs"] = _get_litellm_metadata_from_kwargs + return _globals["get_litellm_metadata_from_kwargs"] + + if name == "map_finish_reason": + # Check if already cached + if "map_finish_reason" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + map_finish_reason as _map_finish_reason, + ) + _globals["map_finish_reason"] = _map_finish_reason + return _globals["map_finish_reason"] + + if name == "process_response_headers": + # Check if already cached + if "process_response_headers" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + process_response_headers as _process_response_headers, + ) + _globals["process_response_headers"] = _process_response_headers + return _globals["process_response_headers"] + + # Lazy load dot_notation_indexing functions to avoid loading at module import time + if name == "delete_nested_value": + # Check if already cached + if "delete_nested_value" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value as _delete_nested_value, + ) + _globals["delete_nested_value"] = _delete_nested_value + return _globals["delete_nested_value"] + + if name == "is_nested_path": + # Check if already cached + if "is_nested_path" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + is_nested_path as _is_nested_path, + ) + _globals["is_nested_path"] = _is_nested_path + return _globals["is_nested_path"] + + # Lazy load get_litellm_params functions to avoid loading at module import time + if name == "_get_base_model_from_litellm_call_metadata": + # Check if already cached + if "_get_base_model_from_litellm_call_metadata" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata as __get_base_model_from_litellm_call_metadata, + ) + _globals["_get_base_model_from_litellm_call_metadata"] = __get_base_model_from_litellm_call_metadata + return _globals["_get_base_model_from_litellm_call_metadata"] + + if name == "get_litellm_params": + # Check if already cached + if "get_litellm_params" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + get_litellm_params as _get_litellm_params, + ) + _globals["get_litellm_params"] = _get_litellm_params + return _globals["get_litellm_params"] + + # Lazy load _ensure_extra_body_is_safe to avoid loading at module import time + if name == "_ensure_extra_body_is_safe": + # Check if already cached + if "_ensure_extra_body_is_safe" not in _globals: + from litellm.litellm_core_utils.llm_request_utils import ( + _ensure_extra_body_is_safe as __ensure_extra_body_is_safe, + ) + _globals["_ensure_extra_body_is_safe"] = __ensure_extra_body_is_safe + return _globals["_ensure_extra_body_is_safe"] + + # Lazy load get_formatted_prompt to avoid loading at module import time + if name == "get_formatted_prompt": + # Check if already cached + if "get_formatted_prompt" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt as _get_formatted_prompt, + ) + _globals["get_formatted_prompt"] = _get_formatted_prompt + return _globals["get_formatted_prompt"] + + # Lazy load get_response_headers to avoid loading at module import time + if name == "get_response_headers": + # Check if already cached + if "get_response_headers" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers as _get_response_headers, + ) + _globals["get_response_headers"] = _get_response_headers + return _globals["get_response_headers"] + + # Lazy load update_response_metadata to avoid loading at module import time + if name == "update_response_metadata": + # Check if already cached + if "update_response_metadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata as _update_response_metadata, + ) + _globals["update_response_metadata"] = _update_response_metadata + return _globals["update_response_metadata"] + + # Lazy load executor to avoid loading at module import time + if name == "executor": + # Check if already cached + if "executor" not in _globals: + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as _executor, + ) + _globals["executor"] = _executor + return _globals["executor"] + + # Lazy load BaseAnthropicMessagesConfig to avoid loading at module import time + if name == "BaseAnthropicMessagesConfig": + # Check if already cached + if "BaseAnthropicMessagesConfig" not in _globals: + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig as _BaseAnthropicMessagesConfig, + ) + _globals["BaseAnthropicMessagesConfig"] = _BaseAnthropicMessagesConfig + return _globals["BaseAnthropicMessagesConfig"] + + # Lazy load BaseAudioTranscriptionConfig to avoid loading at module import time + if name == "BaseAudioTranscriptionConfig": + # Check if already cached + if "BaseAudioTranscriptionConfig" not in _globals: + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig as _BaseAudioTranscriptionConfig, + ) + _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig + return _globals["BaseAudioTranscriptionConfig"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From b6d601c2f02d17ae9409366f24320e6cef70f97b Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 3 Jan 2026 16:16:17 -0800 Subject: [PATCH 248/330] perf(utils): lazy load 15+ unused imports (#18616) - Move BaseBatchesConfig, BaseContainerConfig, BaseEmbeddingConfig, BaseImageEditConfig, BaseImageGenerationConfig, BaseImageVariationConfig, BasePassthroughConfig, BaseRealtimeConfig, BaseRerankConfig, BaseVectorStoreConfig, BaseVectorStoreFilesConfig, BaseVideoConfig to lazy loading - Move ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, RerankResponse to lazy loading - Move ChatCompletionDeltaToolCallChunk, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, LiteLLM_Params to lazy loading - Add type stubs to TYPE_CHECKING block for mypy support - Add lazy loading handlers in __getattr__ method - These imports are not used in utils.py runtime code, only in type annotations (safe with from __future__ import annotations) --- litellm/utils.py | 245 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 216 insertions(+), 29 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 6f4652c8278..df0b2317123 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -149,10 +149,6 @@ def _get_cached_audio_utils(): _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils return _audio_utils_module -from litellm.types.llms.anthropic import ( - ANTHROPIC_API_ONLY_HEADERS, - AnthropicThinkingParam, -) from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -163,7 +159,6 @@ from litellm.types.llms.openai import ( OpenAITextCompletionUserMessage, OpenAIWebSearchOptions, ) -from litellm.types.rerank import RerankResponse from litellm.types.utils import FileTypes # type: ignore from litellm.types.utils import ( OPENAI_RESPONSE_HEADERS, @@ -342,29 +337,41 @@ if TYPE_CHECKING: reset_retry_policy, ) from litellm.secret_managers.main import get_secret + # Type stubs for lazy-loaded config classes and types + from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig + from litellm.llms.base_llm.containers.transformation import BaseContainerConfig + from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig + from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, + ) + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig, + ) + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig + from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig + from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig, + ) + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS, + AnthropicThinkingParam, + ) + from litellm.types.rerank import RerankResponse + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk, + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ) + from litellm.types.router import LiteLLM_Params -from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig -from litellm.llms.base_llm.containers.transformation import BaseContainerConfig -from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.llms.base_llm.image_variations.transformation import ( - BaseImageVariationConfig, -) -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig -from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig -from litellm.llms.base_llm.vector_store_files.transformation import ( - BaseVectorStoreFilesConfig, -) -from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from ._logging import _is_debugging_on, verbose_logger from .caching.caching import ( @@ -392,12 +399,6 @@ from .exceptions import ( UnprocessableEntityError, UnsupportedParamsError, ) -from .types.llms.openai import ( - ChatCompletionDeltaToolCallChunk, - ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, -) -from .types.router import LiteLLM_Params if TYPE_CHECKING: from litellm import MockException @@ -9226,4 +9227,190 @@ def __getattr__(name: str) -> Any: # noqa: PLR0915 _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig return _globals["BaseAudioTranscriptionConfig"] + # Lazy load BaseBatchesConfig to avoid loading at module import time + if name == "BaseBatchesConfig": + # Check if already cached + if "BaseBatchesConfig" not in _globals: + from litellm.llms.base_llm.batches.transformation import ( + BaseBatchesConfig as _BaseBatchesConfig, + ) + _globals["BaseBatchesConfig"] = _BaseBatchesConfig + return _globals["BaseBatchesConfig"] + + # Lazy load BaseContainerConfig to avoid loading at module import time + if name == "BaseContainerConfig": + # Check if already cached + if "BaseContainerConfig" not in _globals: + from litellm.llms.base_llm.containers.transformation import ( + BaseContainerConfig as _BaseContainerConfig, + ) + _globals["BaseContainerConfig"] = _BaseContainerConfig + return _globals["BaseContainerConfig"] + + # Lazy load BaseEmbeddingConfig to avoid loading at module import time + if name == "BaseEmbeddingConfig": + # Check if already cached + if "BaseEmbeddingConfig" not in _globals: + from litellm.llms.base_llm.embedding.transformation import ( + BaseEmbeddingConfig as _BaseEmbeddingConfig, + ) + _globals["BaseEmbeddingConfig"] = _BaseEmbeddingConfig + return _globals["BaseEmbeddingConfig"] + + # Lazy load BaseImageEditConfig to avoid loading at module import time + if name == "BaseImageEditConfig": + # Check if already cached + if "BaseImageEditConfig" not in _globals: + from litellm.llms.base_llm.image_edit.transformation import ( + BaseImageEditConfig as _BaseImageEditConfig, + ) + _globals["BaseImageEditConfig"] = _BaseImageEditConfig + return _globals["BaseImageEditConfig"] + + # Lazy load BaseImageGenerationConfig to avoid loading at module import time + if name == "BaseImageGenerationConfig": + # Check if already cached + if "BaseImageGenerationConfig" not in _globals: + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig as _BaseImageGenerationConfig, + ) + _globals["BaseImageGenerationConfig"] = _BaseImageGenerationConfig + return _globals["BaseImageGenerationConfig"] + + # Lazy load BaseImageVariationConfig to avoid loading at module import time + if name == "BaseImageVariationConfig": + # Check if already cached + if "BaseImageVariationConfig" not in _globals: + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig as _BaseImageVariationConfig, + ) + _globals["BaseImageVariationConfig"] = _BaseImageVariationConfig + return _globals["BaseImageVariationConfig"] + + # Lazy load BasePassthroughConfig to avoid loading at module import time + if name == "BasePassthroughConfig": + # Check if already cached + if "BasePassthroughConfig" not in _globals: + from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig as _BasePassthroughConfig, + ) + _globals["BasePassthroughConfig"] = _BasePassthroughConfig + return _globals["BasePassthroughConfig"] + + # Lazy load BaseRealtimeConfig to avoid loading at module import time + if name == "BaseRealtimeConfig": + # Check if already cached + if "BaseRealtimeConfig" not in _globals: + from litellm.llms.base_llm.realtime.transformation import ( + BaseRealtimeConfig as _BaseRealtimeConfig, + ) + _globals["BaseRealtimeConfig"] = _BaseRealtimeConfig + return _globals["BaseRealtimeConfig"] + + # Lazy load BaseRerankConfig to avoid loading at module import time + if name == "BaseRerankConfig": + # Check if already cached + if "BaseRerankConfig" not in _globals: + from litellm.llms.base_llm.rerank.transformation import ( + BaseRerankConfig as _BaseRerankConfig, + ) + _globals["BaseRerankConfig"] = _BaseRerankConfig + return _globals["BaseRerankConfig"] + + # Lazy load BaseVectorStoreConfig to avoid loading at module import time + if name == "BaseVectorStoreConfig": + # Check if already cached + if "BaseVectorStoreConfig" not in _globals: + from litellm.llms.base_llm.vector_store.transformation import ( + BaseVectorStoreConfig as _BaseVectorStoreConfig, + ) + _globals["BaseVectorStoreConfig"] = _BaseVectorStoreConfig + return _globals["BaseVectorStoreConfig"] + + # Lazy load BaseVectorStoreFilesConfig to avoid loading at module import time + if name == "BaseVectorStoreFilesConfig": + # Check if already cached + if "BaseVectorStoreFilesConfig" not in _globals: + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig as _BaseVectorStoreFilesConfig, + ) + _globals["BaseVectorStoreFilesConfig"] = _BaseVectorStoreFilesConfig + return _globals["BaseVectorStoreFilesConfig"] + + # Lazy load BaseVideoConfig to avoid loading at module import time + if name == "BaseVideoConfig": + # Check if already cached + if "BaseVideoConfig" not in _globals: + from litellm.llms.base_llm.videos.transformation import ( + BaseVideoConfig as _BaseVideoConfig, + ) + _globals["BaseVideoConfig"] = _BaseVideoConfig + return _globals["BaseVideoConfig"] + + # Lazy load ANTHROPIC_API_ONLY_HEADERS to avoid loading at module import time + if name == "ANTHROPIC_API_ONLY_HEADERS": + # Check if already cached + if "ANTHROPIC_API_ONLY_HEADERS" not in _globals: + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS as _ANTHROPIC_API_ONLY_HEADERS, + ) + _globals["ANTHROPIC_API_ONLY_HEADERS"] = _ANTHROPIC_API_ONLY_HEADERS + return _globals["ANTHROPIC_API_ONLY_HEADERS"] + + # Lazy load AnthropicThinkingParam to avoid loading at module import time + if name == "AnthropicThinkingParam": + # Check if already cached + if "AnthropicThinkingParam" not in _globals: + from litellm.types.llms.anthropic import ( + AnthropicThinkingParam as _AnthropicThinkingParam, + ) + _globals["AnthropicThinkingParam"] = _AnthropicThinkingParam + return _globals["AnthropicThinkingParam"] + + # Lazy load RerankResponse to avoid loading at module import time + if name == "RerankResponse": + # Check if already cached + if "RerankResponse" not in _globals: + from litellm.types.rerank import RerankResponse as _RerankResponse + _globals["RerankResponse"] = _RerankResponse + return _globals["RerankResponse"] + + # Lazy load ChatCompletionDeltaToolCallChunk to avoid loading at module import time + if name == "ChatCompletionDeltaToolCallChunk": + # Check if already cached + if "ChatCompletionDeltaToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk as _ChatCompletionDeltaToolCallChunk, + ) + _globals["ChatCompletionDeltaToolCallChunk"] = _ChatCompletionDeltaToolCallChunk + return _globals["ChatCompletionDeltaToolCallChunk"] + + # Lazy load ChatCompletionToolCallChunk to avoid loading at module import time + if name == "ChatCompletionToolCallChunk": + # Check if already cached + if "ChatCompletionToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk as _ChatCompletionToolCallChunk, + ) + _globals["ChatCompletionToolCallChunk"] = _ChatCompletionToolCallChunk + return _globals["ChatCompletionToolCallChunk"] + + # Lazy load ChatCompletionToolCallFunctionChunk to avoid loading at module import time + if name == "ChatCompletionToolCallFunctionChunk": + # Check if already cached + if "ChatCompletionToolCallFunctionChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallFunctionChunk as _ChatCompletionToolCallFunctionChunk, + ) + _globals["ChatCompletionToolCallFunctionChunk"] = _ChatCompletionToolCallFunctionChunk + return _globals["ChatCompletionToolCallFunctionChunk"] + + # Lazy load LiteLLM_Params to avoid loading at module import time + if name == "LiteLLM_Params": + # Check if already cached + if "LiteLLM_Params" not in _globals: + from litellm.types.router import LiteLLM_Params as _LiteLLM_Params + _globals["LiteLLM_Params"] = _LiteLLM_Params + return _globals["LiteLLM_Params"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") From d6296411ec569516a0dabf92ca79b495546146a4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 16:23:10 -0800 Subject: [PATCH 249/330] SSO Settings Loading, deprecate previous flow --- .../AdminSettings/SSOSettings/SSOSettings.tsx | 106 ++++++++++-------- .../SSOSettingsLoadingSkeleton.tsx | 66 +++++++++++ .../src/components/admins.tsx | 8 +- 3 files changed, 130 insertions(+), 50 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx index d339f6d0e36..27ff96af05f 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -10,12 +10,13 @@ import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; import RedactableField from "./RedactableField"; import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants"; const { Title, Text } = Typography; export default function SSOSettings() { - const { data: ssoSettings, refetch } = useSSOSettings(); + const { data: ssoSettings, refetch, isLoading } = useSSOSettings(); const { accessToken } = useAuthorized(); const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false); const [isAddModalVisible, setIsAddModalVisible] = useState(false); @@ -26,27 +27,28 @@ export default function SSOSettings() { Boolean(ssoSettings?.values.generic_client_id); // Determine the SSO provider based on the configuration - let selectedProvider: string | null = null; - if (ssoSettings?.values.google_client_id) { - selectedProvider = "google"; - } else if (ssoSettings?.values.microsoft_client_id) { - selectedProvider = "microsoft"; - } else if (ssoSettings?.values.generic_client_id) { - // Check if it looks like Okta based on endpoints - if ( - ssoSettings.values.generic_authorization_endpoint?.includes("okta") || - ssoSettings.values.generic_authorization_endpoint?.includes("auth0") - ) { - selectedProvider = "okta"; - } else { - selectedProvider = "generic"; + const detectSSOProvider = (values: SSOSettingsValues): string | null => { + if (values.google_client_id) return "google"; + if (values.microsoft_client_id) return "microsoft"; + if (values.generic_client_id) { + // Check if it looks like Okta/Auth0 based on endpoints + if ( + values.generic_authorization_endpoint?.includes("okta") || + values.generic_authorization_endpoint?.includes("auth0") + ) { + return "okta"; + } + return "generic"; } - } + return null; + }; + + const selectedProvider = ssoSettings?.values ? detectSSOProvider(ssoSettings.values) : null; const renderEndpointValue = (value?: string | null) => ( - - {value || Not configured} - + + {value || "-"} + ); const renderSimpleValue = (value?: string | null) => @@ -179,38 +181,44 @@ export default function SSOSettings() { }; return ( - - - {/* Header Section */} -
-
- -
- SSO Configuration - Manage Single Sign-On authentication settings + <> + {isLoading ? ( + + ) : ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ {isSSOConfigured && ( + <> + + + + )} +
-
-
- {isSSOConfigured && ( - <> - - - + {isSSOConfigured ? ( + renderSSOSettings() + ) : ( + setIsAddModalVisible(true)} /> )} -
-
- - {isSSOConfigured ? ( - renderSSOSettings() - ) : ( - setIsAddModalVisible(true)} /> - )} - + + + )} - + ); } diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx new file mode 100644 index 00000000000..59e34f255e3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx @@ -0,0 +1,66 @@ +"use client"; + +import { Card, Descriptions, Skeleton, Space, Typography } from "antd"; +import { Shield } from "lucide-react"; + +const { Title, Text } = Typography; +export default function SSOSettingsLoadingSkeleton() { + const descriptionsConfig = { + column: { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }, + }; + + return ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ + +
+
+ + {/* Descriptions Table Skeleton */} + + {/* Provider Row */} + }> +
+ +
+ + + }> + + + + }> + + + + }> + + + + }> + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx index 6af5a226da6..9de971bcd62 100644 --- a/ui/litellm-dashboard/src/components/admins.tsx +++ b/ui/litellm-dashboard/src/components/admins.tsx @@ -3,7 +3,7 @@ * Use this to avoid sharing master key with others */ import React, { useState, useEffect } from "react"; -import { Typography } from "antd"; +import { Alert, Typography } from "antd"; import { useRouter } from "next/navigation"; import { Button as Button2, Modal, Form, Input } from "antd"; import { Select, SelectItem } from "@tremor/react"; @@ -509,6 +509,12 @@ const AdminPanel: React.FC = ({ ✨ Security Settings +
Date: Sat, 3 Jan 2026 17:25:46 -0800 Subject: [PATCH 250/330] Adding unit testing coverage --- .../Modals/EditSSOSettingsModal.test.tsx | 620 ++++++++++++++++++ .../SSOSettingsLoadingSkeleton.test.tsx | 222 +++++++ .../VectorStoreSelector.test.tsx | 524 +++++++++++++++ .../VectorStoreTable.test.tsx | 415 ++++++++++++ .../src/utils/cookieUtils.test.ts | 72 ++ .../src/utils/proxyUtils.test.ts | 78 +++ .../src/utils/textUtils.test.ts | 21 + 7 files changed, 1952 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx create mode 100644 ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx create mode 100644 ui/litellm-dashboard/src/utils/proxyUtils.test.ts diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..559d837b409 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx @@ -0,0 +1,620 @@ +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, Mock } from "vitest"; +import EditSSOSettingsModal from "./EditSSOSettingsModal"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; + +// Constants +const SSO_PROVIDERS = { + GOOGLE: "google", + MICROSOFT: "microsoft", + OKTA: "okta", + AUTH0: "auth0", + GENERIC: "generic", +} as const; + +const TEST_DATA = { + MODAL_TITLE: "Edit SSO Settings", + MODAL_WIDTH: "800", + SUCCESS_MESSAGE: "SSO settings updated successfully", + ERROR_MESSAGE_PREFIX: "Failed to save SSO settings:", + BUTTON_TEXT: { + CANCEL: "Cancel", + SAVE: "Save", + SAVING: "Saving...", + }, +} as const; + +const TEST_IDS = { + MODAL: "modal", + BUTTON: "button", + BASE_SSO_FORM: "base-sso-form", + TRIGGER_FORM_SUBMIT: "trigger-form-submit", +} as const; + +// Mock form instance +const mockForm = { + resetFields: vi.fn(), + setFieldsValue: vi.fn(), + getFieldsValue: vi.fn(), + submit: vi.fn(), +}; + +// Types +type SSOData = { + values: Record; +} & Record; + +type SSOSettingsHookReturn = { + data: SSOData | null; + isLoading: boolean; + error: any; +}; + +type EditSSOSettingsHookReturn = { + mutateAsync: ReturnType; + isPending: boolean; +}; + +// Test data factories +const createSSOData = (overrides: Record = {}): SSOData => ({ + values: { + user_email: "test@example.com", + ...overrides, + }, +}); + +const createGoogleSSOData = (overrides: Record = {}) => + createSSOData({ + google_client_id: "test-google-id", + google_client_secret: "test-google-secret", + ...overrides, + }); + +const createMicrosoftSSOData = (overrides: Record = {}) => + createSSOData({ + microsoft_client_id: "test-microsoft-id", + microsoft_client_secret: "test-microsoft-secret", + microsoft_tenant: "test-tenant", + ...overrides, + }); + +const createGenericSSOData = (overrides: Record = {}) => + createSSOData({ + generic_client_id: "test-generic-id", + generic_client_secret: "test-generic-secret", + generic_authorization_endpoint: overrides.authorization_endpoint || "https://custom.example.com/oauth", + ...overrides, + }); + +const createRoleMappingsSSOData = (overrides: Record = {}) => + createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + roles: { + proxy_admin: overrides.proxy_admin || ["admin-group"], + proxy_admin_viewer: overrides.proxy_admin_viewer || ["viewer-group"], + internal_user: overrides.internal_user || ["user-group"], + internal_user_viewer: overrides.internal_user_viewer || ["readonly-group"], + }, + }, + ...overrides, + }); + +// Mock utilities +const createMockHooks = (): { + useSSOSettings: SSOSettingsHookReturn; + useEditSSOSettings: EditSSOSettingsHookReturn; +} => ({ + useSSOSettings: { + data: null, + isLoading: false, + error: null, + }, + useEditSSOSettings: { + mutateAsync: vi.fn(), + isPending: false, + }, +}); + +vi.mock("antd", () => ({ + Modal: ({ children, open, title, footer, onCancel, width, ...props }: any) => ( +
+
{children}
+
{footer}
+
+ ), + Button: ({ children, onClick, loading, disabled, ...props }: any) => ( + + ), + Form: { + useForm: () => [mockForm], + }, + Space: ({ children, ...props }: any) => ( +
+ {children} +
+ ), +})); + +vi.mock("./BaseSSOSettingsForm", () => ({ + default: ({ form, onFormSubmit }: any) => ( +
+ +
+ ), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useEditSSOSettings", () => ({ + useEditSSOSettings: vi.fn(), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +vi.mock("../utils", () => ({ + processSSOSettingsPayload: vi.fn(), +})); + +// Test helpers +const setupMocks = ( + overrides: Partial<{ + useSSOSettings: Partial; + useEditSSOSettings: Partial; + }> = {}, +) => { + const defaultMocks = createMockHooks(); + const mocks = { + useSSOSettings: { ...defaultMocks.useSSOSettings, ...overrides.useSSOSettings }, + useEditSSOSettings: { ...defaultMocks.useEditSSOSettings, ...overrides.useEditSSOSettings }, + }; + + (useSSOSettings as Mock).mockReturnValue(mocks.useSSOSettings); + (useEditSSOSettings as Mock).mockReturnValue(mocks.useEditSSOSettings); + + return mocks; +}; + +const renderComponent = (props: Partial> = {}) => { + const defaultProps = { + isVisible: true, + onCancel: vi.fn(), + onSuccess: vi.fn(), + }; + + return { + ...render(), + mockOnCancel: defaultProps.onCancel, + mockOnSuccess: defaultProps.onSuccess, + }; +}; + +const getButtons = () => screen.getAllByTestId(TEST_IDS.BUTTON); +const getCancelButton = () => getButtons()[0]; +const getSaveButton = () => getButtons()[1]; + +describe("EditSSOSettingsModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + setupMocks(); + }); + + describe("Rendering", () => { + it("renders without crashing", () => { + expect(() => renderComponent()).not.toThrow(); + }); + + it("displays modal with correct configuration", () => { + renderComponent(); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "true"); + expect(modal).toHaveAttribute("data-title", TEST_DATA.MODAL_TITLE); + expect(modal).toHaveAttribute("data-width", TEST_DATA.MODAL_WIDTH); + }); + + it("displays modal as closed when not visible", () => { + renderComponent({ isVisible: false }); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "false"); + }); + }); + + describe("Footer Actions", () => { + it("renders cancel and save buttons", () => { + renderComponent(); + + const buttons = getButtons(); + expect(buttons).toHaveLength(2); + expect(buttons[0]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.CANCEL); + expect(buttons[1]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVE); + }); + + it("calls onCancel and resets form when cancel button is clicked", () => { + const { mockOnCancel } = renderComponent(); + + fireEvent.click(getCancelButton()); + + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockOnCancel).toHaveBeenCalled(); + }); + + it("calls form.submit when save button is clicked", () => { + renderComponent(); + + fireEvent.click(getSaveButton()); + + expect(mockForm.submit).toHaveBeenCalled(); + }); + + describe("Loading States", () => { + it("disables cancel button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getCancelButton()).toBeDisabled(); + }); + + it("shows loading state on save button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getSaveButton()).toHaveAttribute("data-loading", "true"); + expect(getSaveButton()).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVING); + }); + }); + }); + + describe("Form Submission", () => { + const formValues = { testField: "testValue" }; + const processedPayload = { processed: "payload" }; + + beforeEach(() => { + (processSSOSettingsPayload as any).mockReturnValue(processedPayload); + }); + + it("processes form values and submits successfully", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalledWith(formValues); + expect(mockMutateAsync).toHaveBeenCalledWith( + processedPayload, + expect.objectContaining({ + onSuccess: expect.any(Function), + onError: expect.any(Function), + }), + ); + }); + + it("shows success notification and calls onSuccess callback", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.success).toHaveBeenCalledWith(TEST_DATA.SUCCESS_MESSAGE); + expect(mockOnSuccess).toHaveBeenCalled(); + }); + + it("handles submission errors gracefully", async () => { + const error = new Error("Submission failed"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue("Parsed error message"); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(parseErrorMessage).toHaveBeenCalledWith(error); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + `${TEST_DATA.ERROR_MESSAGE_PREFIX} Parsed error message`, + ); + }); + }); + + describe("Form Initialization", () => { + describe("Provider Detection", () => { + const testProviderDetection = (testName: string, ssoData: SSOData, expectedProvider: string) => { + it(`detects ${testName} provider`, async () => { + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: expectedProvider, + ...ssoData.values, + }); + }); + }); + }; + + testProviderDetection("Google", createGoogleSSOData(), SSO_PROVIDERS.GOOGLE); + + testProviderDetection("Microsoft", createMicrosoftSSOData(), SSO_PROVIDERS.MICROSOFT); + + testProviderDetection( + "Okta", + createGenericSSOData({ + authorization_endpoint: "https://okta.example.com/oauth2/authorize", + }), + SSO_PROVIDERS.OKTA, + ); + + testProviderDetection( + "Auth0 (detected as Okta)", + createGenericSSOData({ + authorization_endpoint: "https://auth0.example.com/authorize", + }), + SSO_PROVIDERS.OKTA, // Auth0 URLs are detected as Okta provider + ); + + testProviderDetection("generic", createGenericSSOData(), SSO_PROVIDERS.GENERIC); + }); + + describe("Role Mappings", () => { + it("processes role mappings with all roles assigned", async () => { + const ssoData = createRoleMappingsSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "admin-group", + admin_viewer_teams: "viewer-group", + internal_user_teams: "user-group", + internal_viewer_teams: "readonly-group", + }); + }); + }); + + it("handles empty role mapping arrays", async () => { + const ssoData = createRoleMappingsSSOData({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user_viewer: [], + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "user-group", + internal_viewer_teams: "", + }); + }); + }); + }); + + describe("Initialization Guards", () => { + it("resets form before setting values", async () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockForm.setFieldsValue).toHaveBeenCalled(); + }); + }); + + it("skips initialization when modal is not visible", () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent({ isVisible: false }); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + + it("skips initialization when SSO data is unavailable", () => { + setupMocks({ + useSSOSettings: { data: null, isLoading: false, error: null }, + }); + + renderComponent(); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + }); + }); + + describe("Error Handling", () => { + it("handles form submission errors with undefined error message", async () => { + const error = new Error("Network error"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue(undefined); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(`${TEST_DATA.ERROR_MESSAGE_PREFIX} undefined`); + }); + + it("handles form submission with malformed data", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(new Error("Invalid data")); + return Promise.reject(new Error("Invalid data")); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing failed"); + }); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + expect(mockMutateAsync).not.toHaveBeenCalled(); + }); + }); + + describe("Edge Cases", () => { + it("handles role mappings with undefined roles object", async () => { + const ssoData = createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + // roles is undefined + }, + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + }); + }); + }); + + it("handles provider detection with partial SSO data", async () => { + const ssoData = createSSOData({ + // Only has generic fields, no specific provider identifiers + generic_client_id: "test-id", + generic_authorization_endpoint: "https://unknown.provider.com/auth", + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GENERIC, + ...ssoData.values, + }); + }); + }); + + it("handles form submission when processing throws error", async () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing error"); + }); + + renderComponent(); + + expect(() => { + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + }).not.toThrow(); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx new file mode 100644 index 00000000000..fd4fde69588 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx @@ -0,0 +1,222 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; + +// Mock lucide-react icons +vi.mock("lucide-react", () => ({ + Shield: ({ className }: any) =>
, +})); + +// Mock Ant Design components +vi.mock("antd", () => ({ + Card: ({ children, ...props }: any) => ( +
+ {children} +
+ ), + Descriptions: Object.assign( + ({ children, bordered, column, ...props }: any) => ( +
+ {children} +
+ ), + { + Item: ({ children, label, ...props }: any) => ( +
+
{label}
+
{children}
+
+ ), + }, + ), + Typography: { + Title: ({ children, level, ...props }: any) => ( +
+ {children} +
+ ), + Text: ({ children, type, ...props }: any) => ( +
+ {children} +
+ ), + }, + Space: ({ children, direction, size, className, ...props }: any) => ( +
+ {children} +
+ ), + Skeleton: { + Button: ({ active, size, style, ...props }: any) => ( +
+ Button Skeleton +
+ ), + Node: ({ active, style, ...props }: any) => ( +
+ Node Skeleton +
+ ), + }, +})); + +describe("SSOSettingsLoadingSkeleton", () => { + it("should render without crashing", () => { + expect(() => render()).not.toThrow(); + }); + + it("should render Card component", () => { + render(); + expect(screen.getByTestId("card")).toBeInTheDocument(); + }); + + it("should render Space component with correct props", () => { + render(); + const space = screen.getByTestId("space"); + expect(space).toBeInTheDocument(); + expect(space).toHaveAttribute("data-direction", "vertical"); + expect(space).toHaveAttribute("data-size", "large"); + expect(space).toHaveClass("w-full"); + }); + + describe("Header Section", () => { + it("should render Shield icon", () => { + render(); + const shieldIcon = screen.getByTestId("shield-icon"); + expect(shieldIcon).toBeInTheDocument(); + expect(shieldIcon).toHaveClass("w-6 h-6 text-gray-400"); + }); + + it("should render title with correct text and level", () => { + render(); + const title = screen.getByTestId("typography-title"); + expect(title).toBeInTheDocument(); + expect(title).toHaveAttribute("data-level", "3"); + expect(title).toHaveTextContent("SSO Configuration"); + }); + + it("should render subtitle text", () => { + render(); + const text = screen.getByTestId("typography-text"); + expect(text).toBeInTheDocument(); + expect(text).toHaveAttribute("data-type", "secondary"); + expect(text).toHaveTextContent("Manage Single Sign-On authentication settings"); + }); + + it("should render two skeleton buttons with correct styles", () => { + render(); + const buttons = screen.getAllByTestId("skeleton-button"); + expect(buttons).toHaveLength(2); + + // First button + expect(buttons[0]).toHaveAttribute("data-active", "true"); + expect(buttons[0]).toHaveAttribute("data-size", "default"); + expect(buttons[0]).toHaveAttribute("data-style", JSON.stringify({ width: 170, height: 32 })); + + // Second button + expect(buttons[1]).toHaveAttribute("data-active", "true"); + expect(buttons[1]).toHaveAttribute("data-size", "default"); + expect(buttons[1]).toHaveAttribute("data-style", JSON.stringify({ width: 190, height: 32 })); + }); + }); + + describe("Descriptions Table", () => { + it("should render Descriptions component with bordered prop", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + expect(descriptions).toBeInTheDocument(); + expect(descriptions).toHaveAttribute("data-bordered", "true"); + }); + + it("should apply correct column configuration", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + const expectedColumn = { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }; + expect(descriptions).toHaveAttribute("data-column", JSON.stringify(expectedColumn)); + }); + + it("should render exactly 5 description items", () => { + render(); + const items = screen.getAllByTestId("descriptions-item"); + expect(items).toHaveLength(5); + }); + + describe("Description Items Structure", () => { + it("should render exactly 10 skeleton nodes total", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + expect(skeletonNodes).toHaveLength(10); + }); + + it("should render 5 skeleton nodes for labels with width 80", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + const labelNodes = skeletonNodes.filter( + (node) => node.getAttribute("data-style") === JSON.stringify({ width: 80, height: 16 }), + ); + expect(labelNodes).toHaveLength(5); + + labelNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + }); + + it("should render skeleton nodes for content with correct widths", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + // Expected content widths: [100, 200, 250, 180, 220] + const expectedWidths = [100, 200, 250, 180, 220]; + expectedWidths.forEach((width) => { + const contentNode = skeletonNodes.find( + (node) => node.getAttribute("data-style") === JSON.stringify({ width, height: 16 }), + ); + expect(contentNode).toBeInTheDocument(); + expect(contentNode).toHaveAttribute("data-active", "true"); + }); + }); + }); + }); + + describe("Accessibility and Structure", () => { + it("should have proper semantic structure", () => { + render(); + // Card contains Space + const card = screen.getByTestId("card"); + const space = screen.getByTestId("space"); + expect(card).toContainElement(space); + + // Space contains header section and descriptions + const descriptions = screen.getByTestId("descriptions"); + expect(space).toContainElement(descriptions); + }); + + it("should render all skeleton elements as active", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + const skeletonButtons = screen.getAllByTestId("skeleton-button"); + + skeletonNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + + skeletonButtons.forEach((button) => { + expect(button).toHaveAttribute("data-active", "true"); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx new file mode 100644 index 00000000000..8c6b85a53de --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx @@ -0,0 +1,524 @@ +import { render, screen, waitFor, fireEvent } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import VectorStoreSelector from "./VectorStoreSelector"; +import { vectorStoreListCall } from "../networking"; +import { VectorStore } from "./types"; + +// Mock dependencies +const mockVectorStoreListCall = vi.fn(); + +vi.mock("../networking", () => ({ + vectorStoreListCall: (...args: any[]) => mockVectorStoreListCall(...args), +})); + +// Mock antd Select component +vi.mock("antd", () => ({ + Select: vi.fn(), +})); + +// Import the mocked Select +import { Select as MockedSelect } from "antd"; + +// Configure the mock to render a simple div with data attributes +(MockedSelect as any).mockImplementation((props: any) => { + const { + onChange, + value, + placeholder, + loading, + className, + disabled, + options, + mode, + showSearch, + optionFilterProp, + style, + } = props; + + return ( +
{ + // For testing purposes, allow simulating different selection behaviors + // The test can control this by setting data attributes on the element + const testSelection = e.target.getAttribute("data-test-selection"); + if (testSelection && onChange) { + onChange(JSON.parse(testSelection)); + } else if (onChange && options?.length > 0) { + // Default behavior: select first option + onChange([options[0].value]); + } + }} + > + {options?.map((opt: any) => ( +
+ {opt.label} +
+ ))} +
+ ); +}); + +// Test helpers +const mockOnChange = vi.fn(); +const mockAccessToken = "test-token"; + +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "store-1", + custom_llm_provider: "openai", + vector_store_name: "My Store", + vector_store_description: "A test store", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + { + vector_store_id: "store-2", + custom_llm_provider: "azure", + vector_store_name: "Another Store", + vector_store_description: "Another test store", + created_at: "2024-01-02T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + }, + { + vector_store_id: "store-3", + custom_llm_provider: "pg_vector", + // No vector_store_name to test fallback to vector_store_id + vector_store_description: "Store without name", + created_at: "2024-01-03T00:00:00Z", + updated_at: "2024-01-03T00:00:00Z", + }, +]; + +const defaultProps = { + onChange: mockOnChange, + accessToken: mockAccessToken, +}; + +// Helper functions +const renderComponent = (props = {}) => { + return render(); +}; + +const waitForDataFetch = async () => { + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalled(); + }); +}; + +const getSelectElement = () => screen.getByTestId("vector-store-select"); + +const getOptionElements = () => + screen.getAllByTestId(/^vector-store-select/).filter((el) => el.hasAttribute("data-option-value")); + +describe("VectorStoreSelector", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockVectorStoreListCall.mockResolvedValue({ + data: mockVectorStores, + }); + }); + + describe("Rendering", () => { + it("should render the select component", () => { + renderComponent(); + expect(getSelectElement()).toBeInTheDocument(); + }); + + it("should render with default placeholder", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-placeholder", "Select vector stores"); + }); + + it("should render with custom placeholder", () => { + renderComponent({ placeholder: "Choose stores" }); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-placeholder", "Choose stores"); + }); + + it("should apply custom className", () => { + renderComponent({ className: "custom-class" }); + const select = getSelectElement(); + expect(select).toHaveClass("custom-class"); + }); + + it("should render with disabled state", () => { + renderComponent({ disabled: true }); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-disabled", "true"); + }); + + it("should render with enabled state by default", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-disabled", "false"); + }); + + it("should render with multiple mode", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-mode", "multiple"); + }); + + it("should render with showSearch enabled", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-show-search", "true"); + }); + + it("should render with optionFilterProp set to label", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-option-filter-prop", "label"); + }); + + it("should render with full width style", () => { + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveStyle({ width: "100%" }); + }); + }); + + describe("Data fetching", () => { + it("should fetch vector stores on mount when accessToken is provided", async () => { + renderComponent(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith(mockAccessToken); + }); + }); + + it("should not fetch vector stores when accessToken is falsy", () => { + const { rerender } = render(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + + rerender(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + + rerender(); + expect(mockVectorStoreListCall).not.toHaveBeenCalled(); + }); + + it("should fetch vector stores again when accessToken changes", async () => { + const { rerender } = render(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith("token-1"); + }); + + vi.clearAllMocks(); + rerender(); + await waitFor(() => { + expect(mockVectorStoreListCall).toHaveBeenCalledWith("token-2"); + }); + }); + + it("should set loading state while fetching", async () => { + let resolvePromise: (value: any) => void; + const promise = new Promise((resolve) => { + resolvePromise = resolve; + }); + mockVectorStoreListCall.mockReturnValue(promise); + + renderComponent(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "true"); + + resolvePromise!({ data: mockVectorStores }); + await waitFor(() => { + expect(select).toHaveAttribute("data-loading", "false"); + }); + }); + + it("should clear loading state after successful fetch", async () => { + renderComponent(); + await waitForDataFetch(); + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "false"); + }); + + it("should clear loading state after failed fetch", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce(new Error("Network error")); + + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + expect(select).toHaveAttribute("data-loading", "false"); + consoleErrorSpy.mockRestore(); + }); + }); + + describe("Options rendering", () => { + it("should render vector store options after successful fetch", async () => { + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText("My Store (store-1)")).toBeInTheDocument(); + expect(screen.getByText("Another Store (store-2)")).toBeInTheDocument(); + expect(screen.getByText("store-3 (store-3)")).toBeInTheDocument(); + }); + + it("should use vector_store_name when available for label", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toBeInTheDocument(); + expect(option1).toHaveAttribute("data-option-title", "A test store"); + }); + + it("should fallback to vector_store_id when vector_store_name is missing", async () => { + renderComponent(); + await waitForDataFetch(); + + const option3 = screen.getByText("store-3 (store-3)"); + expect(option3).toBeInTheDocument(); + // When vector_store_name is missing, title uses vector_store_description if available, otherwise vector_store_id + expect(option3).toHaveAttribute("data-option-title", "Store without name"); + }); + + it("should use vector_store_description as title when available", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toHaveAttribute("data-option-title", "A test store"); + }); + + it("should fallback to vector_store_id as title when vector_store_description is missing", async () => { + const storesWithoutDescription: VectorStore[] = [ + { + vector_store_id: "store-no-desc", + custom_llm_provider: "openai", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: storesWithoutDescription, + }); + + renderComponent(); + await waitForDataFetch(); + + const option = screen.getByText("store-no-desc (store-no-desc)"); + expect(option).toHaveAttribute("data-option-title", "store-no-desc"); + }); + + it("should use vector_store_id as option value", async () => { + renderComponent(); + await waitForDataFetch(); + + const option1 = screen.getByText("My Store (store-1)"); + expect(option1).toHaveAttribute("data-option-value", "store-1"); + }); + + it("should handle empty vector stores array", async () => { + mockVectorStoreListCall.mockResolvedValueOnce({ + data: [], + }); + + renderComponent(); + await waitForDataFetch(); + + const options = getOptionElements(); + expect(options.length).toBe(0); + }); + + it("should handle response without data property", async () => { + mockVectorStoreListCall.mockResolvedValueOnce({}); + + renderComponent(); + await waitForDataFetch(); + + const options = getOptionElements(); + expect(options.length).toBe(0); + }); + }); + + describe("Value prop", () => { + it("should set initial value when value prop is provided", async () => { + renderComponent({ value: ["store-1", "store-2"] }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBe(JSON.stringify(["store-1", "store-2"])); + }); + + it("should handle empty value array", async () => { + renderComponent({ value: [] }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBe(JSON.stringify([])); + }); + + it("should handle undefined value", async () => { + renderComponent({ value: undefined }); + await waitForDataFetch(); + + const select = getSelectElement(); + const dataValue = select.getAttribute("data-value"); + expect(dataValue).toBeNull(); // undefined value results in no data-value attribute + }); + }); + + describe("onChange callback", () => { + it("should call onChange when selection changes", async () => { + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate selecting store-1 by setting test data attribute + select.setAttribute("data-test-selection", '["store-1"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-1"]); + }); + + it("should call onChange with multiple selected values", async () => { + renderComponent(); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate selecting multiple values + select.setAttribute("data-test-selection", '["store-1", "store-2"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-1", "store-2"]); + }); + + it("should call onChange when deselecting options", async () => { + renderComponent({ value: ["store-1", "store-2"] }); + await waitForDataFetch(); + + const select = getSelectElement(); + // Simulate deselecting store-1 + select.setAttribute("data-test-selection", '["store-2"]'); + fireEvent.click(select); + + expect(mockOnChange).toHaveBeenCalledWith(["store-2"]); + }); + }); + + describe("Error handling", () => { + it("should handle fetch errors gracefully", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const error = new Error("Network error"); + mockVectorStoreListCall.mockRejectedValueOnce(error); + + renderComponent(); + await waitForDataFetch(); + + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching vector stores:", error); + consoleErrorSpy.mockRestore(); + }); + + it("should not crash when fetch throws non-Error", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce("String error"); + + renderComponent(); + await waitForDataFetch(); + + expect(consoleErrorSpy).toHaveBeenCalledWith("Error fetching vector stores:", "String error"); + consoleErrorSpy.mockRestore(); + }); + + it("should continue to work after error", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + mockVectorStoreListCall.mockRejectedValueOnce(new Error("Network error")); + + renderComponent(); + await waitForDataFetch(); + + // Component should still render + expect(getSelectElement()).toBeInTheDocument(); + consoleErrorSpy.mockRestore(); + }); + }); + + describe("Edge cases", () => { + it("should handle vector stores with all optional fields missing", async () => { + const minimalStores: VectorStore[] = [ + { + vector_store_id: "minimal-store", + custom_llm_provider: "openai", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: minimalStores, + }); + + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText("minimal-store (minimal-store)")).toBeInTheDocument(); + const option = screen.getByText("minimal-store (minimal-store)"); + expect(option).toHaveAttribute("data-option-title", "minimal-store"); + }); + + it("should handle very long vector store names", async () => { + const longNameStores: VectorStore[] = [ + { + vector_store_id: "store-long", + custom_llm_provider: "openai", + vector_store_name: "A".repeat(200), + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: longNameStores, + }); + + renderComponent(); + await waitForDataFetch(); + + const expectedLabel = `${"A".repeat(200)} (store-long)`; + expect(screen.getByText(expectedLabel)).toBeInTheDocument(); + }); + + it("should handle special characters in vector store names", async () => { + const specialCharStores: VectorStore[] = [ + { + vector_store_id: "store-special", + custom_llm_provider: "openai", + vector_store_name: 'Store & Co. "Quotes"', + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + mockVectorStoreListCall.mockResolvedValueOnce({ + data: specialCharStores, + }); + + renderComponent(); + await waitForDataFetch(); + + expect(screen.getByText(/Store & Co\. "Quotes"/)).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx new file mode 100644 index 00000000000..16d5d3623eb --- /dev/null +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -0,0 +1,415 @@ +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import VectorStoreTable from "./VectorStoreTable"; +import { VectorStore } from "./types"; + +// Mock dependencies +const mockGetProviderLogoAndName = vi.fn(); +const mockTableIconActionButton = vi.fn(); + +vi.mock("../provider_info_helpers", () => ({ + getProviderLogoAndName: (...args: any[]) => mockGetProviderLogoAndName(...args), +})); + +vi.mock("../common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ + default: (props: any) => { + mockTableIconActionButton(props); + return ( + + ); + }, +})); + +// Mock Tremor components to avoid complex styling issues +vi.mock("@tremor/react", () => ({ + Table: ({ children, ...props }: any) => {children}
, + TableHead: ({ children, ...props }: any) => {children}, + TableBody: ({ children, ...props }: any) => {children}, + TableRow: ({ children, ...props }: any) => {children}, + TableHeaderCell: ({ children, ...props }: any) => {children}, + TableCell: ({ children, ...props }: any) => {children}, +})); + +// Mock antd Tooltip +vi.mock("antd", () => ({ + Tooltip: ({ children, title }: any) => ( +
+ {children} +
+ ), +})); + +// Mock Heroicons +vi.mock("@heroicons/react/outline", () => ({ + ChevronDownIcon: (props: any) =>
, + ChevronUpIcon: (props: any) =>
, + SwitchVerticalIcon: (props: any) =>
, +})); + +// Test data +const mockVectorStores: VectorStore[] = [ + { + vector_store_id: "short-id", + custom_llm_provider: "openai", + vector_store_name: "My OpenAI Store", + vector_store_description: "A store for OpenAI vectors", + created_at: "2024-01-15T10:30:00Z", + updated_at: "2024-01-15T11:00:00Z", + created_by: "user-1", + updated_by: "user-1", + }, + { + vector_store_id: "very-long-vector-store-id-that-should-be-truncated", + custom_llm_provider: "azure", + vector_store_name: undefined, // Test missing name + vector_store_description: "A store for Azure vectors with a very long description that should show a tooltip", + created_at: "2024-01-10T09:15:00Z", + updated_at: "2024-01-12T14:20:00Z", + }, + { + vector_store_id: "store-3", + custom_llm_provider: "pg_vector", + vector_store_name: "PostgreSQL Store", + vector_store_description: undefined, // Test missing description + created_at: "2024-01-05T08:00:00Z", + updated_at: "2024-01-08T16:45:00Z", + }, +]; + +// Mock functions +const mockOnView = vi.fn(); +const mockOnEdit = vi.fn(); +const mockOnDelete = vi.fn(); + +const defaultProps = { + data: mockVectorStores, + onView: mockOnView, + onEdit: mockOnEdit, + onDelete: mockOnDelete, +}; + +// Helper function to render component +const renderComponent = (props = {}) => { + return render(); +}; + +describe("VectorStoreTable", () => { + beforeEach(() => { + vi.clearAllMocks(); + + // Setup default mock returns for getProviderLogoAndName + mockGetProviderLogoAndName.mockImplementation((provider: string) => { + const providerMap: Record = { + openai: { displayName: "OpenAI", logo: "/openai-logo.png" }, + azure: { displayName: "Azure", logo: "/azure-logo.png" }, + pg_vector: { displayName: "PostgreSQL Vector", logo: "/pg-logo.png" }, + }; + return providerMap[provider] || { displayName: provider, logo: "" }; + }); + }); + + describe("Rendering", () => { + it("should render the table with data", () => { + renderComponent(); + expect(screen.getByRole("table")).toBeInTheDocument(); + }); + + it("should render table headers", () => { + renderComponent(); + expect(screen.getByText("Vector Store ID")).toBeInTheDocument(); + expect(screen.getByText("Name")).toBeInTheDocument(); + expect(screen.getByText("Description")).toBeInTheDocument(); + expect(screen.getByText("Provider")).toBeInTheDocument(); + expect(screen.getByText("Created At")).toBeInTheDocument(); + expect(screen.getByText("Updated At")).toBeInTheDocument(); + // Check that we have the expected number of header cells (6 data + 1 actions) + const headers = screen.getAllByRole("columnheader"); + expect(headers).toHaveLength(7); + }); + + it("should render all vector store rows", () => { + renderComponent(); + expect(screen.getAllByRole("row")).toHaveLength(mockVectorStores.length + 1); // +1 for header row + }); + + it("should render empty state when no data", () => { + renderComponent({ data: [] }); + expect(screen.getByText("No vector stores found")).toBeInTheDocument(); + }); + }); + + describe("Vector Store ID Column", () => { + it("should render short vector store IDs fully", () => { + renderComponent(); + expect(screen.getByText("short-id")).toBeInTheDocument(); + }); + + it("should truncate long vector store IDs", () => { + renderComponent(); + // Check that the truncated text is rendered (first 15 chars + ...) + const truncatedText = "very-long-vecto..."; + expect(screen.getByText(truncatedText)).toBeInTheDocument(); + }); + + it("should make vector store ID clickable", async () => { + const user = userEvent.setup(); + renderComponent(); + const idButton = screen.getByText("short-id"); + await user.click(idButton); + expect(mockOnView).toHaveBeenCalledWith("short-id"); + }); + + it("should have correct styling for vector store ID button", () => { + renderComponent(); + const idButton = screen.getByText("short-id").closest("button"); + expect(idButton).toHaveClass("font-mono", "text-blue-500", "bg-blue-50", "hover:bg-blue-100"); + }); + }); + + describe("Name Column", () => { + it("should render vector store name", () => { + renderComponent(); + expect(screen.getByText("My OpenAI Store")).toBeInTheDocument(); + }); + + it("should render fallback for missing name", () => { + renderComponent(); + const fallbackElements = screen.getAllByText("-"); + expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + }); + + it("should wrap name in tooltip", () => { + renderComponent(); + const tooltips = screen.getAllByTestId("tooltip"); + const nameTooltip = tooltips.find((t) => t.getAttribute("data-title") === "My OpenAI Store"); + expect(nameTooltip).toBeInTheDocument(); + }); + }); + + describe("Description Column", () => { + it("should render vector store description", () => { + renderComponent(); + expect(screen.getByText("A store for OpenAI vectors")).toBeInTheDocument(); + }); + + it("should render fallback for missing description", () => { + renderComponent(); + const fallbackElements = screen.getAllByText("-"); + expect(fallbackElements.length).toBe(2); // One for missing name, one for missing description + }); + + it("should wrap description in tooltip", () => { + renderComponent(); + const tooltips = screen.getAllByTestId("tooltip"); + const descTooltip = tooltips.find( + (t) => + t.getAttribute("data-title") === + "A store for Azure vectors with a very long description that should show a tooltip", + ); + expect(descTooltip).toBeInTheDocument(); + }); + }); + + describe("Provider Column", () => { + it("should render provider display name", () => { + renderComponent(); + expect(screen.getByText("OpenAI")).toBeInTheDocument(); + expect(screen.getByText("Azure")).toBeInTheDocument(); + expect(screen.getByText("PostgreSQL Vector")).toBeInTheDocument(); + }); + + it("should render provider logo when available", () => { + renderComponent(); + const logos = screen.getAllByRole("img"); + expect(logos).toHaveLength(3); // All providers have logos in our mock + expect(logos[0]).toHaveAttribute("src", "/openai-logo.png"); + expect(logos[0]).toHaveAttribute("alt", "OpenAI"); + }); + + it("should call getProviderLogoAndName for each provider", () => { + renderComponent(); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("openai"); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("azure"); + expect(mockGetProviderLogoAndName).toHaveBeenCalledWith("pg_vector"); + }); + }); + + describe("Date Columns", () => { + it("should render created at dates", () => { + renderComponent(); + const dateElements = screen.getAllByText(/1\/\d+\/2024/); + expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates + }); + + it("should render updated at dates", () => { + renderComponent(); + const dateElements = screen.getAllByText(/1\/\d+\/2024/); + expect(dateElements.length).toBe(6); // 3 created_at + 3 updated_at dates + }); + }); + + describe("Actions Column", () => { + it("should render edit and delete action buttons for each row", () => { + renderComponent(); + expect(screen.getAllByTestId("action-button-edit")).toHaveLength(mockVectorStores.length); + expect(screen.getAllByTestId("action-button-delete")).toHaveLength(mockVectorStores.length); + }); + + it("should call onEdit when edit button is clicked", async () => { + const user = userEvent.setup(); + renderComponent(); + const editButtons = screen.getAllByTestId("action-button-edit"); + await user.click(editButtons[0]); + expect(mockOnEdit).toHaveBeenCalledWith("short-id"); + }); + + it("should call onDelete when delete button is clicked", async () => { + const user = userEvent.setup(); + renderComponent(); + const deleteButtons = screen.getAllByTestId("action-button-delete"); + await user.click(deleteButtons[0]); + expect(mockOnDelete).toHaveBeenCalledWith("short-id"); + }); + + it("should pass correct props to TableIconActionButton", () => { + renderComponent(); + expect(mockTableIconActionButton).toHaveBeenCalledWith( + expect.objectContaining({ + variant: "Edit", + tooltipText: "Edit vector store", + onClick: expect.any(Function), + }), + ); + expect(mockTableIconActionButton).toHaveBeenCalledWith( + expect.objectContaining({ + variant: "Delete", + tooltipText: "Delete vector store", + onClick: expect.any(Function), + }), + ); + }); + }); + + describe("Sorting", () => { + it("should initialize with created_at descending sort", () => { + renderComponent(); + // The table should initialize with sorting state + expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); + }); + + it("should render sort icons for sortable columns", () => { + renderComponent(); + // Should have sort icons for Created At and Updated At columns + const sortIcons = screen.getAllByTestId(/^chevron-(up|down)$|^switch-vertical$/); + expect(sortIcons.length).toBeGreaterThan(0); + }); + + it("should make header cells clickable for sorting", () => { + renderComponent(); + const headerCells = screen.getAllByRole("columnheader"); + const sortableHeaders = headerCells.filter((cell) => cell.textContent !== ""); + expect(sortableHeaders.length).toBeGreaterThan(0); + }); + + it("should show ascending icon when sorted ascending", () => { + renderComponent(); + // Initially shows descending, but we can test the logic by checking the icons are present + expect(screen.getByTestId("chevron-down")).toBeInTheDocument(); + }); + }); + + describe("Styling and Layout", () => { + it("should apply correct CSS classes to table container", () => { + renderComponent(); + const tableContainer = screen.getByRole("table").parentElement?.parentElement; + expect(tableContainer).toHaveClass("rounded-lg", "custom-border", "relative"); + }); + + it("should apply overflow styling to table wrapper", () => { + renderComponent(); + const tableWrapper = screen.getByRole("table").parentElement; + expect(tableWrapper).toHaveClass("overflow-x-auto"); + }); + + it("should apply sticky styling to actions column", () => { + renderComponent(); + const headerCells = screen.getAllByRole("columnheader"); + const actionsHeader = headerCells[headerCells.length - 1]; + expect(actionsHeader).toHaveClass("sticky", "right-0", "bg-white"); + }); + + it("should apply sticky styling to action cells", () => { + renderComponent(); + const rows = screen.getAllByRole("row").slice(1); // Skip header row + rows.forEach((row) => { + const cells = row.querySelectorAll("td"); + const lastCell = cells[cells.length - 1]; + expect(lastCell).toHaveClass("sticky", "right-0", "bg-white"); + }); + }); + }); + + describe("Table Row Styling", () => { + it("should apply correct height to table rows", () => { + renderComponent(); + const rows = screen.getAllByRole("row").slice(1); // Skip header row + rows.forEach((row) => { + expect(row).toHaveClass("h-8"); + }); + }); + + it("should apply correct cell padding and styling", () => { + renderComponent(); + const cells = screen.getAllByRole("cell"); + cells.forEach((cell) => { + expect(cell).toHaveClass("py-0.5", "max-h-8", "overflow-hidden", "text-ellipsis", "whitespace-nowrap"); + }); + }); + }); + + describe("Empty State", () => { + it("should render single row with centered message when no data", () => { + renderComponent({ data: [] }); + const rows = screen.getAllByRole("row"); + expect(rows).toHaveLength(2); // Header + empty state row + expect(screen.getByText("No vector stores found")).toBeInTheDocument(); + }); + + it("should span all columns in empty state", () => { + renderComponent({ data: [] }); + const emptyCell = screen.getByText("No vector stores found").closest("td"); + expect(emptyCell).toHaveAttribute("colSpan", "7"); // 6 data columns + 1 actions column + }); + }); + + describe("Data Edge Cases", () => { + it("should handle vector stores with minimal data", () => { + const minimalData: VectorStore[] = [ + { + vector_store_id: "minimal", + custom_llm_provider: "test", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ]; + + renderComponent({ data: minimalData }); + expect(screen.getByText("minimal")).toBeInTheDocument(); + expect(screen.getAllByText("-")).toHaveLength(2); // Name and description fallbacks + }); + + it("should handle single vector store", () => { + const singleData = [mockVectorStores[0]]; + renderComponent({ data: singleData }); + expect(screen.getAllByRole("row")).toHaveLength(2); // Header + 1 data row + }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts index 1993819be88..8b066e6a8ea 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -44,6 +44,78 @@ describe("cookieUtils", () => { expect(getCookie("token")).toBeNull(); }); + + it("should return early when window is undefined (server-side rendering)", () => { + const originalWindow = global.window; + const originalDocument = global.document; + + // Mock server-side environment + delete (global as any).window; + delete (global as any).document; + + // This should not throw an error and should return early + expect(() => clearTokenCookies()).not.toThrow(); + + // Restore globals + global.window = originalWindow; + global.document = originalDocument; + }); + + it("should return early when document is undefined (server-side rendering)", () => { + const originalDocument = global.document; + + // Mock server-side environment where document is undefined + delete (global as any).document; + + // This should not throw an error and should return early + expect(() => clearTokenCookies()).not.toThrow(); + + // Restore globals + global.document = originalDocument; + }); + + it("should add current path directory to paths array when different from root and /ui", () => { + // Mock window.location.pathname using vi.stubGlobal + const originalLocation = window.location; + vi.stubGlobal('location', { ...originalLocation, pathname: '/custom/path/page.html' }); + + // Spy on document.cookie to verify the paths being used + const cookieSpy = vi.spyOn(document, 'cookie', 'set'); + + clearTokenCookies(); + + // Verify that cookies were cleared for /custom/path/ path + expect(cookieSpy).toHaveBeenCalledWith( + expect.stringContaining('path=/custom/path/') + ); + + vi.restoreAllMocks(); + }); + + it("should not add current path directory when it's already in paths array", () => { + // Mock window.location.pathname using vi.stubGlobal + const originalLocation = window.location; + vi.stubGlobal('location', { ...originalLocation, pathname: '/' }); + + // Spy on document.cookie to count calls + const cookieSpy = vi.spyOn(document, 'cookie', 'set'); + + clearTokenCookies(); + + // Count how many times each path was used + const rootPathCalls = cookieSpy.mock.calls.filter(call => + call[0].includes('path=/;') || call[0].includes('path=/ ') + ); + const uiPathCalls = cookieSpy.mock.calls.filter(call => + call[0].includes('path=/ui;') || call[0].includes('path=/ui ') + ); + + // Should have calls for root and /ui paths, but not duplicate root + expect(rootPathCalls.length).toBeGreaterThan(0); + expect(uiPathCalls.length).toBeGreaterThan(0); + + vi.restoreAllMocks(); + }); }); describe("getCookie", () => { diff --git a/ui/litellm-dashboard/src/utils/proxyUtils.test.ts b/ui/litellm-dashboard/src/utils/proxyUtils.test.ts new file mode 100644 index 00000000000..37bbd429db9 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/proxyUtils.test.ts @@ -0,0 +1,78 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { fetchProxySettings } from "./proxyUtils"; +import { getProxyUISettings } from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + getProxyUISettings: vi.fn(), +})); + +describe("fetchProxySettings", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("should return null when accessToken is null", async () => { + const result = await fetchProxySettings(null); + + expect(result).toBeNull(); + expect(getProxyUISettings).not.toHaveBeenCalled(); + }); + + it("should return null when accessToken is undefined", async () => { + const result = await fetchProxySettings(undefined as any); + + expect(result).toBeNull(); + expect(getProxyUISettings).not.toHaveBeenCalled(); + }); + + it("should return proxy settings when getProxyUISettings succeeds", async () => { + const mockProxySettings = { someSetting: "value", anotherSetting: 123 }; + const accessToken = "test-token"; + + vi.mocked(getProxyUISettings).mockResolvedValue(mockProxySettings); + + const result = await fetchProxySettings(accessToken); + + expect(result).toEqual(mockProxySettings); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + }); + + it("should return null and log error when getProxyUISettings throws", async () => { + const accessToken = "test-token"; + const mockError = new Error("Network error"); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + vi.mocked(getProxyUISettings).mockRejectedValue(mockError); + + const result = await fetchProxySettings(accessToken); + + expect(result).toBeNull(); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + expect(consoleSpy).toHaveBeenCalledWith("Error fetching proxy settings:", mockError); + + consoleSpy.mockRestore(); + }); + + it("should return null and log error when getProxyUISettings throws a string", async () => { + const accessToken = "test-token"; + const mockError = "String error"; + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + + vi.mocked(getProxyUISettings).mockRejectedValue(mockError); + + const result = await fetchProxySettings(accessToken); + + expect(result).toBeNull(); + expect(getProxyUISettings).toHaveBeenCalledOnce(); + expect(getProxyUISettings).toHaveBeenCalledWith(accessToken); + expect(consoleSpy).toHaveBeenCalledWith("Error fetching proxy settings:", mockError); + + consoleSpy.mockRestore(); + }); +}); \ No newline at end of file diff --git a/ui/litellm-dashboard/src/utils/textUtils.test.ts b/ui/litellm-dashboard/src/utils/textUtils.test.ts index b7c91b9dd33..dfb37ad63b4 100644 --- a/ui/litellm-dashboard/src/utils/textUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/textUtils.test.ts @@ -5,6 +5,15 @@ describe("formatLabel", () => { it("should format label", () => { expect(formatLabel("test_label")).toBe("Test Label"); }); + + it("should return empty string when text is empty string", () => { + expect(formatLabel("")).toBe(""); + }); + + it("should return the same value when text is falsy", () => { + expect(formatLabel(null as any)).toBe(null); + expect(formatLabel(undefined as any)).toBe(undefined); + }); }); describe("truncateString", () => { @@ -26,4 +35,16 @@ describe("formItemValidateJSON", () => { it("should reject with an error message for invalid JSON", async () => { await expect(formItemValidateJSON({}, "invalid JSON")).rejects.toBe("Please enter valid JSON"); }); + + it("should resolve when value is empty string", async () => { + await expect(formItemValidateJSON({}, "")).resolves.toBeUndefined(); + }); + + it("should resolve when value is null", async () => { + await expect(formItemValidateJSON({}, null as any)).resolves.toBeUndefined(); + }); + + it("should resolve when value is undefined", async () => { + await expect(formItemValidateJSON({}, undefined as any)).resolves.toBeUndefined(); + }); }); From 1184db079ec96bee585b08f5e62971534774048f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 17:30:52 -0800 Subject: [PATCH 251/330] fixing tests --- .../Modals/EditSSOSettingsModal.tsx | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx index 297698a7ba0..a731af68ff1 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -88,17 +88,22 @@ const EditSSOSettingsModal: React.FC = ({ isVisible, // Enhanced form submission handler const handleFormSubmit = async (formValues: Record) => { - const payload = processSSOSettingsPayload(formValues); + try { + const payload = processSSOSettingsPayload(formValues); - await mutateAsync(payload, { - onSuccess: () => { - NotificationsManager.success("SSO settings updated successfully"); - onSuccess(); - }, - onError: (error) => { - NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); - }, - }); + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings updated successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + // Handle processing errors gracefully + NotificationsManager.fromBackend("Failed to process SSO settings: " + parseErrorMessage(error)); + } }; const handleCancel = () => { From 816124a40bc8e9d0614e11be10becd5a33352d6c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 17:39:59 -0800 Subject: [PATCH 252/330] Fixign build --- .../vector_store_management/VectorStoreSelector.test.tsx | 6 ++---- .../vector_store_management/VectorStoreTable.test.tsx | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx index 8c6b85a53de..67476b5559d 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreSelector.test.tsx @@ -1,9 +1,7 @@ -import { render, screen, waitFor, fireEvent } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import VectorStoreSelector from "./VectorStoreSelector"; -import { vectorStoreListCall } from "../networking"; import { VectorStore } from "./types"; +import VectorStoreSelector from "./VectorStoreSelector"; // Mock dependencies const mockVectorStoreListCall = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx index 16d5d3623eb..65d15260c4c 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTable.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import VectorStoreTable from "./VectorStoreTable"; From 1112974112ecc4aacf7ca1ab811c7e3cbac48b8f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 3 Jan 2026 19:22:38 -0800 Subject: [PATCH 253/330] Virtual Keys Table Loading State --- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 132 +++++++++++++++++- .../VirtualKeysPage/VirtualKeysTable.tsx | 21 +-- 2 files changed, 142 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 3f55b11769c..cbd3d2c7320 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -1,4 +1,4 @@ -import { screen, waitFor } from "@testing-library/react"; +import { screen, waitFor, fireEvent } from "@testing-library/react"; import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { VirtualKeysTable } from "./VirtualKeysTable"; @@ -264,3 +264,133 @@ it("should show skeleton loaders when isLoading is true", () => { expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); }); + +it("should show 'No keys found' message when filteredKeys is empty", () => { + // Mock empty filteredKeys + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [], + allKeyAliases: [], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + expect(screen.getByText("No keys found")).toBeInTheDocument(); +}); + +it("should handle models with more than 3 entries to trigger expansion UI", () => { + const keyWithManyModels = { + ...mockKey, + models: ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo", "claude-3", "claude-3-5-sonnet"], + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithManyModels], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // This test ensures the ChevronDownIcon import (line 6) is used + // by having a key with > 3 models which triggers the expansion logic + // that uses ChevronDownIcon and ChevronRightIcon + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); +}); + +it("should render table headers correctly", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Check that main headers are rendered (testing the header.isPlaceholder condition path) + expect(screen.getByText("Key ID")).toBeInTheDocument(); + expect(screen.getByText("Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Team Alias")).toBeInTheDocument(); + expect(screen.getByText("Models")).toBeInTheDocument(); + expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); +}); + +it("should handle column resizing hover events", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Find a header cell with data-header-id attribute + const headerCell = document.querySelector("[data-header-id]") as HTMLElement; + + expect(headerCell).toBeInTheDocument(); + + // Check that the resizer element exists within the header + const resizer = headerCell?.querySelector(".resizer") as HTMLElement; + expect(resizer).toBeInTheDocument(); + + // Initially, resizer should have opacity 0 + expect(resizer.style.opacity).toBe("0"); + + // Simulate mouse enter using fireEvent - should set opacity to 0.5 (lines 612-616) + fireEvent.mouseEnter(headerCell); + expect(resizer.style.opacity).toBe("0.5"); + + // Simulate mouse leave using fireEvent - should set opacity back to 0 (lines 618-622) + fireEvent.mouseLeave(headerCell); + expect(resizer.style.opacity).toBe("0"); +}); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index b95d675979c..3bda8ee2f02 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -68,12 +68,13 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo }); const [tablePagination, setTablePagination] = React.useState({ pageIndex: 0, - pageSize: 100, + pageSize: 50, }); const { data: keys, isPending: isLoading, + isFetching, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize); const totalCount = keys?.total_count || 0; @@ -545,8 +546,8 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
- {isLoading ? ( - + {isLoading || isFetching ? ( + ) : ( Showing {rangeLabel} of {totalCount} results @@ -554,32 +555,32 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo )}
- {isLoading ? ( - + {isLoading || isFetching ? ( + ) : ( Page {pageIndex + 1} of {table.getPageCount()} )} - {isLoading ? ( + {isLoading || isFetching ? ( ) : ( )} - {isLoading ? ( + {isLoading || isFetching ? ( ) : ( -
- {isFetching ? ( -
- Loading configuration... -
- ) : Object.keys(discountConfig).length > 0 ? ( - - ) : ( -
- + +
+ Provider Discounts + + Apply percentage-based discounts to reduce costs for specific providers + +
+ + + + + Discounts + Test It + + + +
+
+
- )} -
- - -
- -
- - - - - + {isFetching ? ( +
+ Loading configuration... +
+ ) : Object.keys(discountConfig).length > 0 ? ( + + ) : ( +
+ + + + + No provider discounts configured + + + Click "Add Provider Discount" to get started + +
+ )} +
+ + +
+ +
+ + + + + + )} - {/* Accordion 2: Fee/Price Margin */} - + {/* Accordion 2: Fee/Price Margin - Only for proxy admins */} + {isProxyAdmin && ( + + +
+ Fee/Price Margin + + Add fees or margins to LLM costs for internal billing and cost recovery + +
+ + +
+
+ +
+ {isFetching ? ( +
+ Loading configuration... +
+ ) : Object.keys(marginConfig).length > 0 ? ( + + ) : ( +
+ + + + + No provider margins configured + + + Click "Add Provider Margin" to get started + +
+ )} +
+ + + )} + + {/* Accordion 3: Pricing Calculator - Available to all roles */} +
- Fee/Price Margin + Pricing Calculator - Add fees or margins to LLM costs for internal billing and cost recovery + Estimate LLM costs based on expected token usage and request volume
-
- -
- {isFetching ? ( -
- Loading configuration... -
- ) : Object.keys(marginConfig).length > 0 ? ( - - ) : ( -
- - - - - No provider margins configured - - - Click "Add Provider Margin" to get started - -
- )} +
diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/cost_results.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/cost_results.tsx new file mode 100644 index 00000000000..03d5ca0e518 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/cost_results.tsx @@ -0,0 +1,203 @@ +import React from "react"; +import { Text } from "@tremor/react"; +import { Card, Statistic, Row, Col, Divider, Spin } from "antd"; +import { DollarOutlined, LoadingOutlined } from "@ant-design/icons"; +import { CostEstimateResponse } from "../types"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import ExportDropdown from "./export_dropdown"; + +interface CostResultsProps { + result: CostEstimateResponse | null; + loading: boolean; +} + +const formatCost = (value: number | null | undefined): string => { + if (value === null || value === undefined) return "-"; + if (value === 0) return "$0"; + if (value < 0.0001) return `$${value.toExponential(2)}`; + if (value < 1) return `$${value.toFixed(4)}`; + return `$${formatNumberWithCommas(value, 2, true)}`; +}; + +const formatRequests = (value: number | null | undefined): string => { + if (value === null || value === undefined) return "-"; + return formatNumberWithCommas(value, 0, true); +}; + +const CostResults: React.FC = ({ result, loading }) => { + if (!result && !loading) { + return ( +
+ + Select a model to see cost estimates + +
+ ); + } + + if (loading && !result) { + return ( +
+ } /> + Calculating costs... +
+ ); + } + + if (!result) return null; + + return ( +
+ + +
+
+ Cost Estimate + + Model: {result.model} {result.provider && `(${result.provider})`} + +
+
+ {loading && } size="small" />} + +
+
+ + + + + } + /> + + + + + + + + + 0 ? "#faad14" : undefined, + }} + /> + + + + + {result.daily_cost !== null && ( + + + + } + /> + + + + + + + + + 0 ? "#faad14" : undefined, + }} + /> + + + + )} + + {result.monthly_cost !== null && ( + + + + } + /> + + + + + + + + + 0 ? "#faad14" : undefined, + }} + /> + + + + )} + + {(result.input_cost_per_token || result.output_cost_per_token) && ( +
+ Token Pricing: + {result.input_cost_per_token && ( + Input: ${formatNumberWithCommas(result.input_cost_per_token * 1_000_000, 2)}/1M tokens + )} + {result.input_cost_per_token && result.output_cost_per_token && " | "} + {result.output_cost_per_token && ( + Output: ${formatNumberWithCommas(result.output_cost_per_token * 1_000_000, 2)}/1M tokens + )} +
+ )} +
+ ); +}; + +export default CostResults; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_dropdown.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_dropdown.tsx new file mode 100644 index 00000000000..e8a681021d6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_dropdown.tsx @@ -0,0 +1,71 @@ +import React, { useState, useRef, useEffect } from "react"; +import { Button } from "@tremor/react"; +import { DownloadOutlined, FilePdfOutlined, FileExcelOutlined } from "@ant-design/icons"; +import { CostEstimateResponse } from "../types"; +import { exportToPDF, exportToCSV } from "./export_utils"; + +interface ExportDropdownProps { + result: CostEstimateResponse; +} + +const ExportDropdown: React.FC = ({ result }) => { + const [isOpen, setIsOpen] = useState(false); + const menuRef = useRef(null); + + useEffect(() => { + const handleClickOutside = (event: MouseEvent) => { + if (menuRef.current && !menuRef.current.contains(event.target as Node)) { + setIsOpen(false); + } + }; + + if (isOpen) { + document.addEventListener("mousedown", handleClickOutside); + } + + return () => { + document.removeEventListener("mousedown", handleClickOutside); + }; + }, [isOpen]); + + return ( +
+ + + {isOpen && ( +
+ + +
+ )} +
+ ); +}; + +export default ExportDropdown; + diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_utils.ts b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_utils.ts new file mode 100644 index 00000000000..e02e8288456 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/export_utils.ts @@ -0,0 +1,276 @@ +import { CostEstimateResponse } from "../types"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; + +const formatCostForExport = (value: number | null | undefined): string => { + if (value === null || value === undefined) return "-"; + if (value === 0) return "$0.00"; + if (value < 0.01) return `$${value.toFixed(6)}`; + if (value < 1) return `$${value.toFixed(4)}`; + return `$${formatNumberWithCommas(value, 2)}`; +}; + +const formatRequestsForExport = (value: number | null | undefined): string => { + if (value === null || value === undefined) return "-"; + return formatNumberWithCommas(value, 0); +}; + +export const exportToPDF = (result: CostEstimateResponse): void => { + const printWindow = window.open("", "_blank"); + if (!printWindow) { + alert("Please allow popups to export PDF"); + return; + } + + const html = ` + + + + Cost Estimate Report - ${result.model} + + + +

LLM Cost Estimate Report

+ +
+

Model: ${result.model}

+ ${result.provider ? `

Provider: ${result.provider}

` : ""} +

Input Tokens per Request: ${formatRequestsForExport(result.input_tokens)}

+

Output Tokens per Request: ${formatRequestsForExport(result.output_tokens)}

+ ${result.num_requests_per_day ? `

Requests per Day: ${formatRequestsForExport(result.num_requests_per_day)}

` : ""} + ${result.num_requests_per_month ? `

Requests per Month: ${formatRequestsForExport(result.num_requests_per_month)}

` : ""} +
+ +

Per-Request Cost Breakdown

+ + + + + + + + + + + + + + + + + + + + + +
Cost TypeAmount
Input Cost${formatCostForExport(result.input_cost_per_request)}
Output Cost${formatCostForExport(result.output_cost_per_request)}
Margin/Fee${formatCostForExport(result.margin_cost_per_request)}
Total per Request${formatCostForExport(result.cost_per_request)}
+ + ${result.daily_cost !== null ? ` +

Daily Cost Estimate (${formatRequestsForExport(result.num_requests_per_day)} requests/day)

+ + + + + + + + + + + + + + + + + + + + + +
Cost TypeAmount
Input Cost${formatCostForExport(result.daily_input_cost)}
Output Cost${formatCostForExport(result.daily_output_cost)}
Margin/Fee${formatCostForExport(result.daily_margin_cost)}
Total Daily${formatCostForExport(result.daily_cost)}
+ ` : ""} + + ${result.monthly_cost !== null ? ` +

Monthly Cost Estimate (${formatRequestsForExport(result.num_requests_per_month)} requests/month)

+ + + + + + + + + + + + + + + + + + + + + +
Cost TypeAmount
Input Cost${formatCostForExport(result.monthly_input_cost)}
Output Cost${formatCostForExport(result.monthly_output_cost)}
Margin/Fee${formatCostForExport(result.monthly_margin_cost)}
Total Monthly${formatCostForExport(result.monthly_cost)}
+ ` : ""} + + ${result.input_cost_per_token || result.output_cost_per_token ? ` +

Token Pricing

+ + + + + + ${result.input_cost_per_token ? ` + + + + + ` : ""} + ${result.output_cost_per_token ? ` + + + + + ` : ""} +
Token TypePrice per 1M Tokens
Input Tokens$${(result.input_cost_per_token * 1000000).toFixed(2)}
Output Tokens$${(result.output_cost_per_token * 1000000).toFixed(2)}
+ ` : ""} + +
+

Generated by LiteLLM Pricing Calculator on ${new Date().toLocaleString()}

+
+ + + `; + + printWindow.document.write(html); + printWindow.document.close(); + printWindow.onload = () => { + printWindow.print(); + }; +}; + +export const exportToCSV = (result: CostEstimateResponse): void => { + const rows = [ + ["LLM Cost Estimate Report"], + [""], + ["Configuration"], + ["Model", result.model], + ["Provider", result.provider || "-"], + ["Input Tokens per Request", result.input_tokens.toString()], + ["Output Tokens per Request", result.output_tokens.toString()], + ["Requests per Day", result.num_requests_per_day?.toString() || "-"], + ["Requests per Month", result.num_requests_per_month?.toString() || "-"], + [""], + ["Per-Request Costs"], + ["Input Cost", result.input_cost_per_request.toString()], + ["Output Cost", result.output_cost_per_request.toString()], + ["Margin/Fee", result.margin_cost_per_request.toString()], + ["Total per Request", result.cost_per_request.toString()], + ]; + + if (result.daily_cost !== null) { + rows.push( + [""], + ["Daily Costs"], + ["Daily Input Cost", result.daily_input_cost?.toString() || "-"], + ["Daily Output Cost", result.daily_output_cost?.toString() || "-"], + ["Daily Margin/Fee", result.daily_margin_cost?.toString() || "-"], + ["Total Daily", result.daily_cost.toString()] + ); + } + + if (result.monthly_cost !== null) { + rows.push( + [""], + ["Monthly Costs"], + ["Monthly Input Cost", result.monthly_input_cost?.toString() || "-"], + ["Monthly Output Cost", result.monthly_output_cost?.toString() || "-"], + ["Monthly Margin/Fee", result.monthly_margin_cost?.toString() || "-"], + ["Total Monthly", result.monthly_cost.toString()] + ); + } + + if (result.input_cost_per_token || result.output_cost_per_token) { + rows.push( + [""], + ["Token Pricing (per 1M tokens)"], + ["Input Token Price", result.input_cost_per_token ? `$${(result.input_cost_per_token * 1000000).toFixed(2)}` : "-"], + ["Output Token Price", result.output_cost_per_token ? `$${(result.output_cost_per_token * 1000000).toFixed(2)}` : "-"] + ); + } + + const csv = rows.map(row => row.join(",")).join("\n"); + const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" }); + const url = window.URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = `cost_estimate_${result.model.replace(/\//g, "_")}_${new Date().toISOString().split("T")[0]}.csv`; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + window.URL.revokeObjectURL(url); +}; + diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.tsx new file mode 100644 index 00000000000..fff6475e8a2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.tsx @@ -0,0 +1,31 @@ +import React, { useCallback } from "react"; +import PricingForm from "./pricing_form"; +import CostResults from "./cost_results"; +import { useCostEstimate } from "./use_cost_estimate"; +import { PricingCalculatorProps, PricingFormValues } from "./types"; + +const PricingCalculator: React.FC = ({ + accessToken, + models, +}) => { + const { loading, result, debouncedFetch } = useCostEstimate(accessToken); + + const handleValuesChange = useCallback( + (_changedValues: Partial, allValues: PricingFormValues) => { + if (allValues.model) { + debouncedFetch(allValues); + } + }, + [debouncedFetch] + ); + + return ( +
+ + +
+ ); +}; + +export default PricingCalculator; + diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/pricing_form.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/pricing_form.tsx new file mode 100644 index 00000000000..c91a19e5516 --- /dev/null +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/pricing_form.tsx @@ -0,0 +1,104 @@ +import React from "react"; +import { Form, InputNumber, Select, Row, Col } from "antd"; +import { PricingFormValues } from "./types"; + +interface PricingFormProps { + models: string[]; + onValuesChange: (changedValues: Partial, allValues: PricingFormValues) => void; +} + +const PricingForm: React.FC = ({ models, onValuesChange }) => { + return ( +
+ + + + handleEntryChange(record.id, "model", value)} + optionFilterProp="label" + filterOption={(input, option) => + String(option?.label ?? "").toLowerCase().includes(input.toLowerCase()) + } + options={models.map((model) => ({ + value: model, + label: model, + }))} + style={{ width: "100%" }} + size="small" + /> + ), + }, + { + title: "Input Tokens", + dataIndex: "input_tokens", + key: "input_tokens", + width: "18%", + render: (_: number, record: ModelEntry) => ( + handleEntryChange(record.id, "input_tokens", value ?? 0)} + style={{ width: "100%" }} + size="small" + formatter={(value) => `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",")} + /> + ), + }, + { + title: "Output Tokens", + dataIndex: "output_tokens", + key: "output_tokens", + width: "18%", + render: (_: number, record: ModelEntry) => ( + handleEntryChange(record.id, "output_tokens", value ?? 0)} + style={{ width: "100%" }} + size="small" + formatter={(value) => `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",")} + /> + ), + }, + { + title: `Requests/${timePeriod === "day" ? "Day" : "Month"}`, + dataIndex: timePeriod === "day" ? "num_requests_per_day" : "num_requests_per_month", + key: "num_requests", + width: "20%", + render: (_: number | undefined, record: ModelEntry) => ( + + handleEntryChange( + record.id, + timePeriod === "day" ? "num_requests_per_day" : "num_requests_per_month", + value ?? undefined + ) + } + style={{ width: "100%" }} + size="small" + placeholder="-" + formatter={(value) => (value ? `${value}`.replace(/\B(?=(\d{3})+(?!\d))/g, ",") : "")} + /> + ), + }, + { + title: "", + key: "actions", + width: 50, + render: (_: unknown, record: ModelEntry) => ( +