From 3f3efea3016bb52d93601d32f326fd02e8ac4313 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Sep 2025 15:12:21 -0700 Subject: [PATCH 001/474] test(test_gemini.py): add additional testing for additionalproperties case --- litellm/types/llms/vertex_ai.py | 1 - tests/llm_translation/test_gemini.py | 79 ++++++++++++++----- .../test_amazing_vertex_completion.py | 50 +++++++++--- 3 files changed, 96 insertions(+), 34 deletions(-) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 2687b79f727..f17a284ddfc 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -113,7 +113,6 @@ class Schema(TypedDict, total=False): pattern: str example: Any anyOf: List["Schema"] - additionalProperties: Any class FunctionDeclaration(TypedDict, total=False): diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index 47e3aaa8143..122429a87a5 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -269,7 +269,11 @@ def test_gemini_image_generation(): assert len(response.choices[0].message.images) > 0 assert response.choices[0].message.images[0]["image_url"] is not None assert response.choices[0].message.images[0]["image_url"]["url"] is not None - assert response.choices[0].message.images[0]["image_url"]["url"].startswith("data:image/png;base64,") + assert ( + response.choices[0] + .message.images[0]["image_url"]["url"] + .startswith("data:image/png;base64,") + ) def test_gemini_thinking(): @@ -661,7 +665,8 @@ def test_system_message_with_no_user_message(): assert response is not None assert response.choices[0].message.content is not None - + + def get_current_weather(location, unit="fahrenheit"): """Get the current weather in a given location""" if "tokyo" in location.lower(): @@ -778,9 +783,9 @@ def test_gemini_reasoning_effort_minimal(): # Test with different Gemini models to verify model-specific mapping test_cases = [ - ("gemini/gemini-2.5-flash", 1), # Flash: minimum 1 token - ("gemini/gemini-2.5-pro", 128), # Pro: minimum 128 tokens - ("gemini/gemini-2.5-flash-lite", 512), # Flash-Lite: minimum 512 tokens + ("gemini/gemini-2.5-flash", 1), # Flash: minimum 1 token + ("gemini/gemini-2.5-pro", 128), # Pro: minimum 128 tokens + ("gemini/gemini-2.5-flash-lite", 512), # Flash-Lite: minimum 512 tokens ] for model, expected_min_budget in test_cases: @@ -793,24 +798,32 @@ def test_gemini_reasoning_effort_minimal(): "reasoning_effort": "minimal", }, ) - + # Verify that the thinking config is set correctly request_body = raw_request["raw_request_body"] - assert "generationConfig" in request_body, f"Model {model} should have generationConfig" - + assert ( + "generationConfig" in request_body + ), f"Model {model} should have generationConfig" + generation_config = request_body["generationConfig"] - assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig" - + assert ( + "thinkingConfig" in generation_config + ), f"Model {model} should have thinkingConfig" + thinking_config = generation_config["thinkingConfig"] - assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget" - + assert ( + "thinkingBudget" in thinking_config + ), f"Model {model} should have thinkingBudget" + actual_budget = thinking_config["thinkingBudget"] - assert actual_budget == expected_min_budget, \ - f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" - + assert ( + actual_budget == expected_min_budget + ), f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + # Verify that includeThoughts is True for minimal reasoning effort - assert thinking_config.get("includeThoughts", True), \ - f"Model {model} should have includeThoughts=True for minimal reasoning effort" + assert thinking_config.get( + "includeThoughts", True + ), f"Model {model} should have includeThoughts=True for minimal reasoning effort" # Test with unknown model (should use generic fallback) try: @@ -822,15 +835,41 @@ def test_gemini_reasoning_effort_minimal(): "reasoning_effort": "minimal", }, ) - + request_body = raw_request["raw_request_body"] generation_config = request_body["generationConfig"] thinking_config = generation_config["thinkingConfig"] # Should use generic fallback (128 tokens) - assert thinking_config["thinkingBudget"] == 128, \ - "Unknown model should use generic fallback of 128 tokens" + assert ( + thinking_config["thinkingBudget"] == 128 + ), "Unknown model should use generic fallback of 128 tokens" except Exception as e: # If return_raw_request doesn't work for unknown models, that's okay # The important part is that our known models work correctly print(f"Note: Unknown model test skipped due to: {e}") pass + + +def test_gemini_additional_properties_bug(): + # Simple tool with additionalProperties (simulating the TypedDict issue) + tools = [ + { + "type": "function", + "function": { + "name": "test_tool", + "description": "Test tool", + "parameters": { + "type": "object", + "properties": {"param1": {"type": "string"}}, + # This causes the error - any non-False value + "additionalProperties": True, # Could also be None, {}, etc. + }, + }, + } + ] + + messages = [{"role": "user", "content": "Test message"}] + + response = litellm.completion( + model="gemini/gemini-2.5-flash", messages=messages, tools=tools + ) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index a27fe738c7f..d20f54bdd34 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -397,7 +397,7 @@ async def test_async_vertexai_response(): | litellm.vertex_text_models | litellm.vertex_code_text_models ) - + test_models = random.sample(list(test_models), 1) test_models += list(litellm.vertex_language_models) # always test gemini-pro for model in test_models: @@ -504,7 +504,6 @@ async def test_async_vertexai_streaming_response(): pytest.fail(f"An exception occurred: {e}") - @pytest.mark.parametrize("load_pdf", [False]) # True, @pytest.mark.flaky(retries=3, delay=1) def test_completion_function_plus_pdf(load_pdf): @@ -547,6 +546,7 @@ def test_completion_function_plus_pdf(load_pdf): except Exception as e: pytest.fail("Got={}".format(str(e))) + def encode_image(image_path): import base64 @@ -910,7 +910,10 @@ async def test_partner_models_httpx(model, region, sync_mode): [ ("vertex_ai/meta/llama-4-scout-17b-16e-instruct-maas", "us-east5"), ("vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", "us-south1"), - ("vertex_ai/mistral-large-2411", "us-central1"), # critical - we had this issue: https://github.com/BerriAI/litellm/issues/13888 + ( + "vertex_ai/mistral-large-2411", + "us-central1", + ), # critical - we had this issue: https://github.com/BerriAI/litellm/issues/13888 ("vertex_ai/openai/gpt-oss-20b-maas", "us-central1"), ], ) @@ -3827,7 +3830,7 @@ def test_vertex_ai_gemini_audio_ogg(): @pytest.mark.asyncio async def test_vertex_ai_deepseek(): """Test that deepseek models use the correct v1 API endpoint instead of v1beta1.""" - #load_vertex_ai_credentials() + # load_vertex_ai_credentials() litellm._turn_on_debug() from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -3840,21 +3843,17 @@ async def test_vertex_ai_deepseek(): { "message": { "role": "assistant", - "content": "Hello! How can I help you today?" + "content": "Hello! How can I help you today?", }, "index": 0, - "finish_reason": "stop" + "finish_reason": "stop", } ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 - }, - "model": "deepseek-ai/deepseek-r1-0528-maas" + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + "model": "deepseek-ai/deepseek-r1-0528-maas", } mock_response.status_code = 200 - + with patch.object(client, "post", return_value=mock_response) as mock_post: response = await acompletion( model="vertex_ai/deepseek-ai/deepseek-r1-0528-maas", @@ -3900,3 +3899,28 @@ def test_gemini_grounding_on_streaming(): vertex_ai_grounding_metadata_shows_up = True print(chunk) assert vertex_ai_grounding_metadata_shows_up + + +def test_gemini_additional_properties_bug(): + # Simple tool with additionalProperties (simulating the TypedDict issue) + tools = [ + { + "type": "function", + "function": { + "name": "test_tool", + "description": "Test tool", + "parameters": { + "type": "object", + "properties": {"param1": {"type": "string"}}, + # This causes the error - any non-False value + "additionalProperties": True, # Could also be None, {}, etc. + }, + }, + } + ] + + messages = [{"role": "user", "content": "Test message"}] + + response = litellm.completion( + model="gemini/gemini-2.5-flash", messages=messages, tools=tools + ) From 84a7329dba2838d6de6e9c1e6a4a7c03dfb9076d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Sep 2025 16:04:06 -0700 Subject: [PATCH 002/474] fix(secret_managers/get_azure_Ad_token_providers.py): infer credential type from env var don't default to ClientSecretCredential unless present in env var --- litellm/llms/azure/common_utils.py | 4 ++- litellm/proxy/_new_secret_config.yaml | 8 +++++ .../get_azure_ad_token_provider.py | 36 +++++++++++++++++-- 3 files changed, 44 insertions(+), 4 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 09b1888e04d..b36375e4168 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -561,7 +561,9 @@ class BaseAzureLLM(BaseOpenAILLM): "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" ) try: - azure_ad_token_provider = get_azure_ad_token_provider(azure_scope=scope) + azure_ad_token_provider = get_azure_ad_token_provider( + azure_scope=scope, + ) except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") if api_version is None: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c785dd05c40..be2a8303864 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -7,3 +7,11 @@ model_list: - model_name: wildcard_models/* litellm_params: model: openai/* + - model_name: gpt-4o + litellm_params: + model: azure/gpt-4o + api_base: https://cog-cuda-atg-sage-eastus2.openai.azure.com + api_version: 2025-04-01-preview + +litellm_settings: + enable_azure_ad_token_refresh: true \ No newline at end of file diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index e73c8f8d7a2..184d959b964 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,11 +1,36 @@ import os from typing import Any, Callable, Optional, Union +from litellm._logging import verbose_logger from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, ) +def infer_credential_type_from_environment() -> AzureCredentialType: + if ( + os.environ.get("AZURE_CLIENT_ID") + and os.environ.get("AZURE_CLIENT_SECRET") + and os.environ.get("AZURE_TENANT_ID") + ): + return AzureCredentialType.ClientSecretCredential + elif os.environ.get("AZURE_CLIENT_ID"): + return AzureCredentialType.ManagedIdentityCredential + elif ( + os.environ.get("AZURE_CLIENT_ID") + and os.environ.get("AZURE_TENANT_ID") + and os.environ.get("AZURE_CERTIFICATE_PATH") + and os.environ.get("AZURE_CERTIFICATE_PASSWORD") + ): + return AzureCredentialType.CertificateCredential + elif os.environ.get("AZURE_CERTIFICATE_PASSWORD"): + return AzureCredentialType.CertificateCredential + elif os.environ.get("AZURE_CERTIFICATE_PATH"): + return AzureCredentialType.CertificateCredential + else: + return AzureCredentialType.DefaultAzureCredential + + def get_azure_ad_token_provider( azure_scope: Optional[str] = None, azure_credential: Optional[AzureCredentialType] = None, @@ -42,9 +67,14 @@ def get_azure_ad_token_provider( ) cred: str = ( - azure_credential.value if azure_credential else None - or os.environ.get("AZURE_CREDENTIAL", AzureCredentialType.ClientSecretCredential) - or AzureCredentialType.ClientSecretCredential + azure_credential.value + if azure_credential + else None + or os.environ.get("AZURE_CREDENTIAL") + or infer_credential_type_from_environment() + ) + verbose_logger.info( + f"For Azure AD Token Provider, choosing credential type: {cred}" ) credential: Optional[ Union[ From d0732f55b3954da5a6853447a99fb8ae5a38edbe Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Sep 2025 16:07:32 -0700 Subject: [PATCH 003/474] test(test_get_azure_ad_token_provider.py): add unit test to ensure default azure credentials used in the right context --- .../test_get_azure_ad_token_provider.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py index 85e55a5c30d..f02f59cccc0 100644 --- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py @@ -214,3 +214,32 @@ class TestGetAzureAdTokenProvider: # Test that the returned callable works token = result() assert token == "mock-certificate-token" + + @patch.dict(os.environ, {}, clear=True) # Clear all environment variables + @patch("azure.identity.get_bearer_token_provider") + @patch("azure.identity.DefaultAzureCredential") + def test_get_azure_ad_token_provider_defaults_to_default_azure_credential( + self, mock_default_azure_credential, mock_get_bearer_token_provider + ): + """Test get_azure_ad_token_provider defaults to DefaultAzureCredential when no credentials are present.""" + # Mock the Azure identity credential instance + mock_credential_instance = MagicMock() + mock_default_azure_credential.return_value = mock_credential_instance + + # Mock the bearer token provider + mock_token_provider = MagicMock(return_value="mock-default-token") + mock_get_bearer_token_provider.return_value = mock_token_provider + + # Call the function + result = get_azure_ad_token_provider() + + # Assertions + assert callable(result) + mock_default_azure_credential.assert_called_once_with() + mock_get_bearer_token_provider.assert_called_once_with( + mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ) + + # Test that the returned callable works + token = result() + assert token == "mock-default-token" From 54e71bd0775adeab4df7015e334dd7ac3cbc551b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Sep 2025 16:10:13 -0700 Subject: [PATCH 004/474] fix(common_utils.py): add helpful message --- litellm/llms/azure/common_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index b36375e4168..7c744298fbd 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -365,6 +365,11 @@ def get_azure_ad_token( azure_ad_token_provider = get_azure_ad_token_provider(azure_scope=scope) except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") + except Exception as e: + verbose_logger.error( + f"Error calling Azure AD token provider: {str(e)}. Follow docs - https://docs.litellm.ai/docs/providers/azure/#azure-ad-token-refresh---defaultazurecredential" + ) + raise e ######################################################### # If litellm.enable_azure_ad_token_refresh is True and no other token provider is available, From 62622ef4b296efaf80dba825ce65e78bd3dac24e Mon Sep 17 00:00:00 2001 From: daily-kim Date: Sun, 21 Sep 2025 10:44:47 +0000 Subject: [PATCH 005/474] fix: update authorization header to use 'Bearer' instead of 'bearer' --- litellm/llms/cohere/common_utils.py | 4 ++-- litellm/llms/cohere/rerank/transformation.py | 2 +- litellm/llms/infinity/rerank/transformation.py | 2 +- .../proxy/pass_through_endpoints/pass_through_endpoints.py | 2 +- tests/local_testing/test_pass_through_endpoints.py | 6 +++--- tests/test_passthrough_endpoints.py | 2 +- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/litellm/llms/cohere/common_utils.py b/litellm/llms/cohere/common_utils.py index 6dbe52d575e..d194d9556b6 100644 --- a/litellm/llms/cohere/common_utils.py +++ b/litellm/llms/cohere/common_utils.py @@ -31,7 +31,7 @@ def validate_environment( "Request-Source": "unspecified:litellm", "accept": "application/json", "content-type": "application/json", - "Authorization": "bearer $CO_API_KEY" + "Authorization": "Bearer $CO_API_KEY" } """ headers.update( @@ -42,7 +42,7 @@ def validate_environment( } ) if api_key: - headers["Authorization"] = f"bearer {api_key}" + headers["Authorization"] = f"Bearer {api_key}" return headers diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 5371b9a4b61..4683ea479f7 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -86,7 +86,7 @@ class CohereRerankConfig(BaseRerankConfig): ) default_headers = { - "Authorization": f"bearer {api_key}", + "Authorization": f"Bearer {api_key}", "accept": "application/json", "content-type": "application/json", } diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 4b75fa121b2..408595cd979 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -49,7 +49,7 @@ class InfinityRerankConfig(CohereRerankConfig): ) default_headers = { - "Authorization": f"bearer {api_key}", + "Authorization": f"Bearer {api_key}", "accept": "application/json", "content-type": "application/json", } diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a1f43d0ca50..f55816edf9e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -76,7 +76,7 @@ async def set_env_variables_in_header(custom_headers: Optional[dict]) -> Optiona example header can be - {"Authorization": "bearer os.environ/COHERE_API_KEY"} + {"Authorization": "Bearer os.environ/COHERE_API_KEY"} """ if custom_headers is None: return None diff --git a/tests/local_testing/test_pass_through_endpoints.py b/tests/local_testing/test_pass_through_endpoints.py index 6cc6a66007f..9836b262efa 100644 --- a/tests/local_testing/test_pass_through_endpoints.py +++ b/tests/local_testing/test_pass_through_endpoints.py @@ -117,7 +117,7 @@ async def test_pass_through_endpoint_rerank(client): { "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] @@ -193,7 +193,7 @@ async def test_pass_through_endpoint_rpm_limit( "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", "auth": auth, - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] @@ -293,7 +293,7 @@ async def test_pass_through_endpoint_sequential_rpm_limit( "path": "/v1/rerank", "target": "https://api.cohere.com/v1/rerank", "auth": auth, - "headers": {"Authorization": f"bearer {_cohere_api_key}"}, + "headers": {"Authorization": f"Bearer {_cohere_api_key}"}, } ] diff --git a/tests/test_passthrough_endpoints.py b/tests/test_passthrough_endpoints.py index a66c94c5836..47ac7511aa1 100644 --- a/tests/test_passthrough_endpoints.py +++ b/tests/test_passthrough_endpoints.py @@ -17,7 +17,7 @@ dotenv.load_dotenv() async def cohere_rerank(session): url = "http://localhost:4000/v1/rerank" headers = { - "Authorization": f"bearer {os.getenv('COHERE_API_KEY')}", + "Authorization": f"Bearer {os.getenv('COHERE_API_KEY')}", "Content-Type": "application/json", "Accept": "application/json", } From a2793bdb5760cb75f4eb944da9d7f0813c128dbb Mon Sep 17 00:00:00 2001 From: Shagun Bansal Date: Tue, 23 Sep 2025 18:49:37 +0530 Subject: [PATCH 006/474] #14404 BugFix - Add support for Azure AD token-based authorization in image generation request headers definition for Azure --- litellm/images/main.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/litellm/images/main.py b/litellm/images/main.py index 2a8b62bce24..0e11d9d3e56 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -314,6 +314,15 @@ def image_generation( # noqa: PLR0915 "Content-Type": "application/json", "api-key": api_key, } + + if api_key is None and azure_ad_token_provider is not None: + azure_ad_token = azure_ad_token_provider() + if azure_ad_token: + default_headers.pop( + "api-key", None + ) + default_headers["Authorization"] = f"Bearer {azure_ad_token}" + for k, v in default_headers.items(): if k not in headers: headers[k] = v From 250e13ea928c79e8d6882f1014914d2899337d16 Mon Sep 17 00:00:00 2001 From: Shagun Bansal Date: Wed, 24 Sep 2025 16:23:20 +0530 Subject: [PATCH 007/474] Revert "#14404 BugFix - Add support for Azure AD token-based authorization in image generation request headers definition for Azure" This reverts commit a2793bdb5760cb75f4eb944da9d7f0813c128dbb. --- litellm/images/main.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index 0e11d9d3e56..2a8b62bce24 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -314,15 +314,6 @@ def image_generation( # noqa: PLR0915 "Content-Type": "application/json", "api-key": api_key, } - - if api_key is None and azure_ad_token_provider is not None: - azure_ad_token = azure_ad_token_provider() - if azure_ad_token: - default_headers.pop( - "api-key", None - ) - default_headers["Authorization"] = f"Bearer {azure_ad_token}" - for k, v in default_headers.items(): if k not in headers: headers[k] = v From 589c83b88b45b53dd5c6bce05ace5dd1245188e8 Mon Sep 17 00:00:00 2001 From: Shagun Bansal Date: Wed, 24 Sep 2025 16:25:44 +0530 Subject: [PATCH 008/474] #14404 BugFix - Add support for Azure AD token-based authorization in image generation request headers definition for Azure --- litellm/llms/azure/azure.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 5ee9065f5e1..f41a9bea0c9 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1117,6 +1117,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): status_code=422, message="max retries must be an int" ) + if api_key is None and azure_ad_token_provider is not None: + azure_ad_token = azure_ad_token_provider() + if azure_ad_token: + headers.pop( + "api-key", None + ) + headers["Authorization"] = f"Bearer {azure_ad_token}" + # init AzureOpenAI Client azure_client_params: Dict[str, Any] = self.initialize_azure_sdk_client( litellm_params=litellm_params or {}, From 6c95bd926f291c0312ea12fae3702465a419f3d7 Mon Sep 17 00:00:00 2001 From: Toy-97 Date: Fri, 26 Sep 2025 20:11:26 +0800 Subject: [PATCH 009/474] update: DeepInfra model data refresh [2025-09-26] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Added models: deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus Removed models: deepinfra/zai-org/GLM-4.5-Air Modified models: deepinfra/NousResearch/Hermes-3-Llama-3.1-70B: - input_cost_per_token: 1.2e-07 → 3e-07 deepinfra/Qwen/Qwen3-32B: - output_cost_per_token: 3e-07 → 2.8e-07 deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct: - max_tokens: 4096 → 262144 - max_output_tokens: 4096 → 262144 - max_input_tokens: 4096 → 262144 deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking: - max_tokens: 4096 → 262144 - max_output_tokens: 4096 → 262144 - max_input_tokens: 4096 → 262144 deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507: - input_cost_per_token: 1.3e-07 → 9e-08 deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct: - input_cost_per_token: 2.3e-07 → 4e-07 deepinfra/google/gemini-2.5-flash: - output_cost_per_token: 1.75e-06 → 2.5e-06 - input_cost_per_token: 2.1e-07 → 3e-07 deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo: - output_cost_per_token: 2e-08 → 3e-08 - input_cost_per_token: 1.5e-08 → 2e-08 deepinfra/meta-llama/Llama-3.2-3B-Instruct: - output_cost_per_token: 2.4e-08 → 2e-08 - input_cost_per_token: 1.2e-08 → 2e-08 deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo: - input_cost_per_token: 2e-08 → 4e-08 deepinfra/openai/gpt-oss-120b: - input_cost_per_token: 9e-08 → 5e-08 deepinfra/google/gemini-2.5-pro: - output_cost_per_token: 7e-06 → 1e-05 - input_cost_per_token: 8.75e-07 → 1.25e-06 deepinfra/NousResearch/Hermes-3-Llama-3.1-405B: - output_cost_per_token: 8e-07 → 1e-06 - input_cost_per_token: 7e-07 → 1e-06 deepinfra/Qwen/Qwen3-235B-A22B: - output_cost_per_token: 6e-07 → 5.4e-07 - input_cost_per_token: 1.3e-07 → 1.8e-07 deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct: - output_cost_per_token: 3e-07 → 6e-07 - input_cost_per_token: 1.2e-07 → 6e-07 deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo: - output_cost_per_token: 1.2e-07 → 3.9e-07 - input_cost_per_token: 3.8e-08 → 1.3e-07 deepinfra/deepseek-ai/DeepSeek-V3-0324: - input_cost_per_token: 2.8e-07 → 2.5e-07 - cache_read_input_token_cost: 2.24e-07 → None deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506: - output_cost_per_token: 1e-07 → 2e-07 - input_cost_per_token: 5e-08 → 7.5e-08 deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507: - output_cost_per_token: 6e-07 → 2.9e-06 - input_cost_per_token: 1.3e-07 → 3e-07 deepinfra/zai-org/GLM-4.5: - output_cost_per_token: 2e-06 → 1.6e-06 - input_cost_per_token: 5.5e-07 → 4e-07 deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1: - output_cost_per_token: 2.4e-07 → 4e-07 - input_cost_per_token: 8e-08 → 4e-07 deepinfra/openai/gpt-oss-20b: - output_cost_per_token: 1.6e-07 → 1.5e-07 deepinfra/google/gemma-3-27b-it: - output_cost_per_token: 1.7e-07 → 1.6e-07 --- model_prices_and_context_window.json | 567 +++++++++++++++------------ 1 file changed, 308 insertions(+), 259 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 755318159b0..235c19af410 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6575,629 +6575,678 @@ ] }, "deepinfra/Gryphe/MythoMax-L2-13b": { - "input_cost_per_token": 7.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 4096, "max_input_tokens": 4096, "max_output_tokens": 4096, - "max_tokens": 4096, + "input_cost_per_token": 8e-08, + "output_cost_per_token": 9e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7.2e-08, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-405B": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1e-06, + "output_cost_per_token": 1e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 8e-07, "supports_tool_choice": true }, "deepinfra/NousResearch/Hermes-3-Llama-3.1-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.8e-07, "supports_tool_choice": false }, "deepinfra/Qwen/QwQ-32B": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-72B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1.2e-07, "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen2.5-7B-Instruct": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Qwen/Qwen2.5-VL-32B-Instruct": { - "input_cost_per_token": 2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, - "mode": "chat", + "input_cost_per_token": 2e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-14B": { - "input_cost_per_token": 6e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 6e-08, "output_cost_per_token": 2.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1.8e-07, + "output_cost_per_token": 5.4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Instruct-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 9e-08, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-235B-A22B-Thinking-2507": { - "input_cost_per_token": 1.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.9e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 6e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-30B-A3B": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 2.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-32B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct": { - "input_cost_per_token": 4e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Qwen/Qwen3-Coder-480B-A35B-Instruct-Turbo": { - "cache_read_input_token_cost": 2.4e-07, - "input_cost_per_token": 3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 262144, "max_input_tokens": 262144, "max_output_tokens": 262144, - "max_tokens": 262144, - "mode": "chat", + "input_cost_per_token": 2.9e-07, "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 1.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/Sao10K/L3-8B-Lunaris-v1-Turbo": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.1-70B-Euryale-v2.2": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/Sao10K/L3.3-70B-Euryale-v2.3": { - "input_cost_per_token": 6.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 6.5e-07, "output_cost_per_token": 7.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/allenai/olmOCR-7B-0725-FP8": { - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.5e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/anthropic/claude-3-7-sonnet-latest": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "cache_read_input_token_cost": 3.3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-opus": { - "input_cost_per_token": 1.65e-05, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 1.65e-05, "output_cost_per_token": 8.25e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/anthropic/claude-4-sonnet": { - "input_cost_per_token": 3.3e-06, - "litellm_provider": "deepinfra", + "max_tokens": 200000, "max_input_tokens": 200000, "max_output_tokens": 200000, - "max_tokens": 200000, - "mode": "chat", + "input_cost_per_token": 3.3e-06, "output_cost_per_token": 1.65e-05, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1": { - "input_cost_per_token": 7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 7e-07, "output_cost_per_token": 2.4e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528": { - "cache_read_input_token_cost": 4e-07, - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2.15e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-0528-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 4e-07, "supports_tool_choice": false }, "deepinfra/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { - "input_cost_per_token": 7.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 2.7e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.5e-07, "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-R1-Turbo": { - "input_cost_per_token": 1e-06, - "litellm_provider": "deepinfra", + "max_tokens": 40960, "max_input_tokens": 40960, "max_output_tokens": 40960, - "max_tokens": 40960, - "mode": "chat", + "input_cost_per_token": 1e-06, "output_cost_per_token": 3e-06, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3": { - "input_cost_per_token": 3.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 3.8e-07, "output_cost_per_token": 8.9e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3-0324": { - "cache_read_input_token_cost": 2.24e-07, - "input_cost_per_token": 2.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.5e-07, "output_cost_per_token": 8.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/deepseek-ai/DeepSeek-V3.1": { - "cache_read_input_token_cost": 2.16e-07, - "input_cost_per_token": 2.7e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1e-06, - "supports_reasoning": true, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/deepseek-ai/DeepSeek-V3.1-Terminus": { + "max_tokens": 163840, + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "input_cost_per_token": 2.7e-07, + "output_cost_per_token": 1e-06, + "cache_read_input_token_cost": 2.16e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.0-flash-001": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-flash": { - "input_cost_per_token": 2.1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 2.5e-06, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.75e-06, "supports_tool_choice": true }, "deepinfra/google/gemini-2.5-pro": { - "input_cost_per_token": 8.75e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1000000, "max_input_tokens": 1000000, "max_output_tokens": 1000000, - "max_tokens": 1000000, + "input_cost_per_token": 1.25e-06, + "output_cost_per_token": 1e-05, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 7e-06, "supports_tool_choice": true }, "deepinfra/google/gemma-3-12b-it": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 1e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/google/gemma-3-27b-it": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.7e-07, "supports_tool_choice": true }, "deepinfra/google/gemma-3-4b-it": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.2-11B-Vision-Instruct": { - "input_cost_per_token": 4.9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4.9e-08, "output_cost_per_token": 4.9e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-3.2-3B-Instruct": { - "input_cost_per_token": 1.2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 2e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-08, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2.3e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 3.8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 3.9e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.2e-07, "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { - "input_cost_per_token": 1.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 1048576, "max_input_tokens": 1048576, "max_output_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", + "input_cost_per_token": 1.5e-07, "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-4-Scout-17B-16E-Instruct": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 327680, "max_input_tokens": 327680, "max_output_tokens": 327680, - "max_tokens": 327680, - "mode": "chat", + "input_cost_per_token": 8e-08, "output_cost_per_token": 3e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Llama-Guard-3-8B": { - "input_cost_per_token": 5.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5.5e-08, "output_cost_per_token": 5.5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Llama-Guard-4-12B": { - "input_cost_per_token": 1.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 163840, "max_input_tokens": 163840, "max_output_tokens": 163840, - "max_tokens": 163840, - "mode": "chat", + "input_cost_per_token": 1.8e-07, "output_cost_per_token": 1.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/meta-llama/Meta-Llama-3-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 8192, "max_input_tokens": 8192, "max_output_tokens": 8192, - "max_tokens": 8192, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 6e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct": { - "input_cost_per_token": 2.3e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 4e-07, "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { - "input_cost_per_token": 1e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 1e-07, "output_cost_per_token": 2.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct": { - "input_cost_per_token": 3e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 3e-08, "output_cost_per_token": 5e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { - "input_cost_per_token": 1.5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 2e-08, + "output_cost_per_token": 3e-08, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2e-08, "supports_tool_choice": true }, "deepinfra/microsoft/WizardLM-2-8x22B": { - "input_cost_per_token": 4.8e-07, - "litellm_provider": "deepinfra", + "max_tokens": 65536, "max_input_tokens": 65536, "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", + "input_cost_per_token": 4.8e-07, "output_cost_per_token": 4.8e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": false }, "deepinfra/microsoft/phi-4": { - "input_cost_per_token": 7e-08, - "litellm_provider": "deepinfra", + "max_tokens": 16384, "max_input_tokens": 16384, "max_output_tokens": 16384, - "max_tokens": 16384, - "mode": "chat", + "input_cost_per_token": 7e-08, "output_cost_per_token": 1.4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Nemo-Instruct-2407": { - "input_cost_per_token": 2e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 2e-08, "output_cost_per_token": 4e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-24B-Instruct-2501": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 8e-08, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/mistralai/Mistral-Small-3.2-24B-Instruct-2506": { - "input_cost_per_token": 5e-08, - "litellm_provider": "deepinfra", + "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "max_tokens": 128000, + "input_cost_per_token": 7.5e-08, + "output_cost_per_token": 2e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1e-07, "supports_tool_choice": true }, "deepinfra/mistralai/Mixtral-8x7B-Instruct-v0.1": { - "input_cost_per_token": 8e-08, - "litellm_provider": "deepinfra", + "max_tokens": 32768, "max_input_tokens": 32768, "max_output_tokens": 32768, - "max_tokens": 32768, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 2.4e-07, "supports_tool_choice": true }, "deepinfra/moonshotai/Kimi-K2-Instruct": { - "input_cost_per_token": 5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-07, "output_cost_per_token": 2e-06, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/moonshotai/Kimi-K2-Instruct-0905": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/nvidia/Llama-3.1-Nemotron-70B-Instruct": { - "input_cost_per_token": 1.2e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/Llama-3.3-Nemotron-Super-49B-v1.5": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_tool_choice": true + }, + "deepinfra/nvidia/NVIDIA-Nemotron-Nano-9B-v2": { + "max_tokens": 131072, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.6e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 3e-07, "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-120b": { - "input_cost_per_token": 9e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", + "input_cost_per_token": 5e-08, "output_cost_per_token": 4.5e-07, + "litellm_provider": "deepinfra", + "mode": "chat", "supports_tool_choice": true }, "deepinfra/openai/gpt-oss-20b": { - "input_cost_per_token": 4e-08, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, + "input_cost_per_token": 4e-08, + "output_cost_per_token": 1.5e-07, + "litellm_provider": "deepinfra", "mode": "chat", - "output_cost_per_token": 1.6e-07, "supports_tool_choice": true }, "deepinfra/zai-org/GLM-4.5": { - "input_cost_per_token": 5.5e-07, - "litellm_provider": "deepinfra", + "max_tokens": 131072, "max_input_tokens": 131072, "max_output_tokens": 131072, - "max_tokens": 131072, - "mode": "chat", - "output_cost_per_token": 2e-06, - "supports_tool_choice": true - }, - "deepinfra/zai-org/GLM-4.5-Air": { - "input_cost_per_token": 2e-07, + "input_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, "litellm_provider": "deepinfra", - "max_input_tokens": 131072, - "max_output_tokens": 131072, - "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.1e-06, "supports_tool_choice": true }, "deepseek/deepseek-chat": { From 3c99d2236a69dc7846bca53b9d6ce0ece6c99295 Mon Sep 17 00:00:00 2001 From: TobiMayr Date: Sun, 28 Sep 2025 17:13:40 +0100 Subject: [PATCH 010/474] feature/add max requests env var --- docs/my-website/docs/proxy/deploy.md | 19 +++++++++ docs/my-website/docs/proxy/prod.md | 10 +++++ litellm/proxy/proxy_cli.py | 17 ++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 45 ++++++++++++++++++++++ 4 files changed, 91 insertions(+) diff --git a/docs/my-website/docs/proxy/deploy.md b/docs/my-website/docs/proxy/deploy.md index 6a11d069fb0..dc2da22684b 100644 --- a/docs/my-website/docs/proxy/deploy.md +++ b/docs/my-website/docs/proxy/deploy.md @@ -715,6 +715,25 @@ docker run ghcr.io/berriai/litellm:main-stable ``` +### Restart Workers After N Requests + +Use this to mitigate memory growth by recycling workers after a fixed number of requests. When set, each worker restarts after completing the specified number of requests. Defaults to disabled when unset. + +Usage Examples: + +```shell showLineNumbers title="docker run (CLI flag)" +docker run ghcr.io/berriai/litellm:main-stable \ + --max_requests_before_restart 10000 +``` + +Or set via environment variable: + +```shell showLineNumbers title="Environment Variable" +export MAX_REQUESTS_BEFORE_RESTART=10000 +docker run ghcr.io/berriai/litellm:main-stable +``` + + ### 5. config.yaml file on s3, GCS Bucket Object/url Use this if you cannot mount a config file on your deployment service (example - AWS Fargate, Railway etc) diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index a45474f39e8..2858132c8e8 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -71,6 +71,16 @@ Use this Docker `CMD`. This will start the proxy with 1 Uvicorn Async Worker CMD ["--port", "4000", "--config", "./proxy_server_config.yaml"] ``` +> Optional: If you observe gradual memory growth under sustained load, consider recycling workers after a fixed number of requests to mitigate leaks. Set this via CLI or environment variable: + +```shell +# CLI +CMD ["--port", "4000", "--config", "./proxy_server_config.yaml", "--max_requests_before_restart", "10000"] + +# or ENV (for deployment manifests / containers) +export MAX_REQUESTS_BEFORE_RESTART=10000 +``` + ## 4. Use Redis 'port','host', 'password'. NOT 'redis_url' diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 867a395d627..21e2c54ec74 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -185,6 +185,7 @@ class ProxyInitializationHelpers: num_workers: int, ssl_certfile_path: str, ssl_keyfile_path: str, + max_requests_before_restart: Optional[int] = None, ): """ Run litellm with `gunicorn` @@ -265,6 +266,10 @@ class ProxyInitializationHelpers: "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s', } + # Optional: recycle workers after N requests to mitigate memory growth + if max_requests_before_restart is not None: + gunicorn_options["max_requests"] = max_requests_before_restart + if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" # noqa @@ -486,6 +491,13 @@ class ProxyInitializationHelpers: help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)", envvar="KEEPALIVE_TIMEOUT", ) +@click.option( + "--max_requests_before_restart", + default=None, + type=int, + help="Restart worker after this many requests (uvicorn: limit_max_requests, gunicorn: max_requests)", + envvar="MAX_REQUESTS_BEFORE_RESTART", +) def run_server( # noqa: PLR0915 host, port, @@ -524,6 +536,7 @@ def run_server( # noqa: PLR0915 use_prisma_db_push: bool, skip_server_startup, keepalive_timeout, + max_requests_before_restart, ): args = locals() if local: @@ -813,6 +826,9 @@ def run_server( # noqa: PLR0915 log_config=log_config, keepalive_timeout=keepalive_timeout, ) + # Optional: recycle uvicorn workers after N requests + if max_requests_before_restart is not None: + uvicorn_args["limit_max_requests"] = max_requests_before_restart if run_gunicorn is False and run_hypercorn is False: if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa @@ -837,6 +853,7 @@ def run_server( # noqa: PLR0915 num_workers=num_workers, ssl_certfile_path=ssl_certfile_path, ssl_keyfile_path=ssl_keyfile_path, + max_requests_before_restart=max_requests_before_restart, ) elif run_hypercorn is True: ProxyInitializationHelpers._init_hypercorn_server( diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4235e5d3adb..90d958e711d 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -314,6 +314,51 @@ class TestProxyInitializationHelpers: call_args = mock_uvicorn_run.call_args assert call_args[1]["timeout_keep_alive"] == 30 + @patch("uvicorn.run") + @patch("builtins.print") + def test_max_requests_before_restart_flag(self, mock_print, mock_uvicorn_run): + """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + + mock_app = MagicMock() + mock_proxy_config = MagicMock() + mock_key_mgmt = MagicMock() + mock_save_worker_config = MagicMock() + + with patch.dict( + "sys.modules", + { + "proxy_server": MagicMock( + app=mock_app, + ProxyConfig=mock_proxy_config, + KeyManagementSettings=mock_key_mgmt, + save_worker_config=mock_save_worker_config, + ) + }, + ), patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args: + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, ["--local", "--max_requests_before_restart", "123"] + ) + + assert result.exit_code == 0 + mock_uvicorn_run.assert_called_once() + + # Check that uvicorn.run was called with limit_max_requests parameter + call_args = mock_uvicorn_run.call_args + assert call_args[1]["limit_max_requests"] == 123 + @patch.dict(os.environ, {}, clear=True) def test_construct_database_url_from_env_vars(self): """Test the construct_database_url_from_env_vars function with various scenarios""" From aeae6cffe48d4ab1a67d474e9e1e96c20fd68429 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Speglich?= Date: Mon, 29 Sep 2025 20:42:42 -0300 Subject: [PATCH 011/474] oci: drop params automatically and add DEDICATED Support --- litellm/llms/oci/chat/transformation.py | 33 +++++++++++++++++++------ litellm/types/llms/oci.py | 6 ++--- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 6755cab22e0..72a044e014a 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -207,9 +207,9 @@ class OCIChatConfig(BaseConfig): alias = open_ai_to_oci_param_map.get(key) if alias is False: - if drop_params: - continue - + # Workaround for mypy issue + #if drop_params: + continue raise Exception(f"param `{key}` is not supported on OCI") if alias is None: @@ -450,12 +450,26 @@ class OCIChatConfig(BaseConfig): "Cohere models are not yet supported in the litellm OCI chat completion endpoint. Use the Cohere API directly." ) else: - data = OCICompletionPayload( - compartmentId=oci_compartment_id, - servingMode=OCIServingMode( + oci_serving_mode = optional_params.get("oci_serving_mode", "ON_DEMAND") + if oci_serving_mode not in ["ON_DEMAND", "DEDICATED"]: + raise Exception( + "kwarg `oci_serving_mode` must be either 'ON_DEMAND' or 'DEDICATED'" + ) + + if oci_serving_mode == "DEDICATED": + servingMode = OCIServingMode( + servingType="DEDICATED", + endpointId=model, + ) + else: + servingMode = OCIServingMode( servingType="ON_DEMAND", modelId=model, - ), + ) + + data = OCICompletionPayload( + compartmentId=oci_compartment_id, + servingMode=servingMode, chatRequest=OCIChatRequestPayload( apiFormat=vendor.value, messages=adapt_messages_to_generic_oci_standard(messages), @@ -601,6 +615,11 @@ class OCIChatConfig(BaseConfig): if "stream" in data: del data["stream"] + stops = data.get("chatRequest", {}).get("stop") + if stops and len(stops) > 8: + # mantém apenas os 8 primeiros + data["chatRequest"]["stop"] = stops[:8] + if client is None or isinstance(client, HTTPHandler): client = get_async_httpx_client(llm_provider=LlmProviders.BYTEZ, params={}) diff --git a/litellm/types/llms/oci.py b/litellm/types/llms/oci.py index 75d13192c50..56fd61ad7e6 100644 --- a/litellm/types/llms/oci.py +++ b/litellm/types/llms/oci.py @@ -100,8 +100,8 @@ class OCIServingMode(BaseModel): """Defines the serving mode and the model to be used.""" servingType: str - modelId: str - + endpointId: Optional[str] = None + modelId: Optional[str] = None class OCICompletionPayload(BaseModel): """Pydantic model for the complete OCI chat request body.""" @@ -129,7 +129,7 @@ class OCIPromptTokensDetails(BaseModel): class OCIResponseUsage(BaseModel): """Token usage in the OCI response.""" - + promptTokens: int completionTokens: int totalTokens: int From 9810c05a9962690300da742595714ae246598f25 Mon Sep 17 00:00:00 2001 From: Kowyo Date: Tue, 30 Sep 2025 04:21:36 +0000 Subject: [PATCH 012/474] fix(ollama/chat): 'think' param handling --- litellm/llms/ollama/chat/transformation.py | 3 +++ litellm/llms/ollama/completion/transformation.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 3b755e79330..dfd53a5e2ba 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -281,6 +281,7 @@ class OllamaChatConfig(BaseConfig): stream = optional_params.pop("stream", False) format = optional_params.pop("format", None) keep_alive = optional_params.pop("keep_alive", None) + think = optional_params.pop("think", None) function_name = optional_params.pop("function_name", None) litellm_params["function_name"] = function_name tools = optional_params.pop("tools", None) @@ -344,6 +345,8 @@ class OllamaChatConfig(BaseConfig): data["tools"] = tools if keep_alive is not None: data["keep_alive"] = keep_alive + if think is not None: + data["think"] = think return data diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 981a987ec91..662affcb278 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -412,6 +412,7 @@ class OllamaConfig(BaseConfig): stream = optional_params.pop("stream", False) format = optional_params.pop("format", None) images = optional_params.pop("images", None) + think = optional_params.pop("think", None) data = { "model": model, "prompt": ollama_prompt, @@ -425,6 +426,8 @@ class OllamaConfig(BaseConfig): data["images"] = [ _convert_image(convert_to_ollama_image(image)) for image in images ] + if think is not None: + data["think"] = think return data From eecb2ba77a4c7a82fa2cf128a98d9dbd1f8c97b7 Mon Sep 17 00:00:00 2001 From: Kowyo Date: Tue, 30 Sep 2025 12:46:22 +0000 Subject: [PATCH 013/474] fix: add 'think' parameter handling in ollama_chat.py --- litellm/llms/ollama_chat.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/llms/ollama_chat.py b/litellm/llms/ollama_chat.py index 082312d28f2..e186636de99 100644 --- a/litellm/llms/ollama_chat.py +++ b/litellm/llms/ollama_chat.py @@ -59,6 +59,7 @@ def get_ollama_response( # noqa: PLR0915 stream = optional_params.pop("stream", False) format = optional_params.pop("format", None) keep_alive = optional_params.pop("keep_alive", None) + think = optional_params.pop("think", None) function_name = optional_params.pop("function_name", None) tools = optional_params.pop("tools", None) @@ -98,6 +99,8 @@ def get_ollama_response( # noqa: PLR0915 data["tools"] = tools if keep_alive is not None: data["keep_alive"] = keep_alive + if think is not None: + data["think"] = think ## LOGGING logging_obj.pre_call( input=None, From 7ec6a3684a6d64737f7f1bd8c7d0edbe84851022 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Speglich?= Date: Tue, 30 Sep 2025 10:46:14 -0300 Subject: [PATCH 014/474] oci: undo stop crop --- litellm/llms/oci/chat/transformation.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 72a044e014a..a3ee3d58f07 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -615,11 +615,6 @@ class OCIChatConfig(BaseConfig): if "stream" in data: del data["stream"] - stops = data.get("chatRequest", {}).get("stop") - if stops and len(stops) > 8: - # mantém apenas os 8 primeiros - data["chatRequest"]["stop"] = stops[:8] - if client is None or isinstance(client, HTTPHandler): client = get_async_httpx_client(llm_provider=LlmProviders.BYTEZ, params={}) From 60063202517cc9748cfb2530e93313133a3e5753 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 12:50:21 -0700 Subject: [PATCH 015/474] fix(proxy/utils.py): run guardrails before running other logging hooks on "async_post_call_success_hook" Closes LIT-1152 --- litellm/proxy/utils.py | 58 ++++++++++++++++++++++-------------------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5b11c25b2bf..d3d2972abfa 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1395,9 +1395,12 @@ class ProxyLogging: 3. /image/generation 4. /files """ + from litellm.types.guardrails import GuardrailEventHooks - for callback in litellm.callbacks: - try: + guardrail_callbacks: List[CustomGuardrail] = [] + other_callbacks: List[CustomLogger] = [] + try: + for callback in litellm.callbacks: _callback: Optional[CustomLogger] = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -1407,36 +1410,37 @@ class ProxyLogging: _callback = callback # type: ignore if _callback is not None: + if isinstance(_callback, CustomGuardrail): + guardrail_callbacks.append(_callback) + else: + other_callbacks.append(_callback) ############## Handle Guardrails ######################################## ############################################################################# - if isinstance(callback, CustomGuardrail): - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks - if ( - callback.should_run_guardrail( - data=data, event_type=GuardrailEventHooks.post_call - ) - is not True - ): - continue + for callback in guardrail_callbacks: + # Main - V2 Guardrails implementation + if ( + callback.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + continue - await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ) - ############ Handle CustomLogger ############################### - ################################################################# - elif isinstance(_callback, CustomLogger): - await _callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, - ) - except Exception as e: - raise e + ############ Handle CustomLogger ############################### + ################################################################# + for callback in other_callbacks: + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, data=data, response=response + ) + except Exception as e: + raise e return response async def async_post_call_streaming_hook( From 6ca7752381fe8ac6aae08985879d59efdedae6fc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 15:46:30 -0700 Subject: [PATCH 016/474] fix(prometheus.py): don't require metadata labels to be set for all requests add a default value if metadata label not set --- .../integrations/prometheus.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index d3b0aefb86f..a42f9b642d9 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -1649,9 +1649,22 @@ class PrometheusLogger(CustomLogger): api_base: Optional[str], api_provider: str, ): - self.litellm_deployment_state.labels( - litellm_model_name, model_id, api_base, api_provider - ).set(state) + """ + Set the deployment state. + """ + ### get labels + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_state" + ), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + api_provider=api_provider, + ), + ) + self.litellm_deployment_state.labels(**_labels).set(state) def set_deployment_healthy( self, @@ -2230,6 +2243,10 @@ def prometheus_label_factory( for key, value in enum_values.custom_metadata_labels.items(): if key in supported_enum_labels: filtered_labels[key] = value + else: + filtered_labels[key] = ( + "None" # this happens for dynamically added metadata labels + ) # Add custom tags if configured if enum_values.tags is not None: From d6800ee706194aeaff40bbacc653b74033a33586 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 17:02:40 -0700 Subject: [PATCH 017/474] feat(prometheus.py): initial working commit of passing team/key metadata as prometheus metrics Closes LIT-1006 --- .../integrations/prometheus.py | 12 +++-- litellm/litellm_core_utils/litellm_logging.py | 2 + litellm/proxy/_new_secret_config.yaml | 33 +++--------- litellm/proxy/_types.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 52 ++++++++++++++++++- litellm/types/utils.py | 31 ++++++++--- 6 files changed, 93 insertions(+), 38 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index a42f9b642d9..b472ccbb357 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -794,9 +794,16 @@ class PrometheusLogger(CustomLogger): output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] - _requester_metadata = standard_logging_payload["metadata"].get( + _requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get( "requester_metadata" ) + user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[ + "metadata" + ].get("user_api_key_auth_metadata") + combined_metadata: Dict[str, Any] = { + **(_requester_metadata if _requester_metadata else {}), + **(user_api_key_auth_metadata if user_api_key_auth_metadata else {}), + } if standard_logging_payload is not None and isinstance( standard_logging_payload, dict ): @@ -828,8 +835,7 @@ class PrometheusLogger(CustomLogger): exception_status=None, exception_class=None, custom_metadata_labels=get_custom_labels_from_metadata( - metadata=standard_logging_payload["metadata"].get("requester_metadata") - or {} + metadata=combined_metadata ), route=standard_logging_payload["metadata"].get( "user_api_key_request_route" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 24449e1bd0f..46e363c865d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4019,6 +4019,7 @@ class StandardLoggingPayloadSetup: usage_object=usage_object, requester_custom_headers=None, cold_storage_object_key=None, + user_api_key_auth_metadata=None, ) if isinstance(metadata, dict): # Filter the metadata dictionary to include only the specified keys @@ -4685,6 +4686,7 @@ def get_standard_logging_metadata( requester_custom_headers=None, user_api_key_request_route=None, cold_storage_object_key=None, + user_api_key_auth_metadata=None, ) if isinstance(metadata, dict): # Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 804cf2cf2cf..96d9ea0cc31 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,30 +1,9 @@ model_list: - - model_name: byok-fixed-gpt-4o-mini + - model_name: openai/gpt-4o litellm_params: - model: openai/gpt-4o-mini - api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5" - api_key: dummy - - model_name: "byok-wildcard/*" - litellm_params: - model: openai/* - - model_name: xai-grok-3 - litellm_params: - model: xai/grok-3 - - model_name: hosted_vllm/whisper-v3 - litellm_params: - model: hosted_vllm/whisper-v3 - api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5" - api_key: dummy - -mcp_servers: - github_mcp: - url: "https://api.githubcopilot.com/mcp" - auth_type: oauth2 - authorization_url: https://github.com/login/oauth/authorize - token_url: https://github.com/login/oauth/access_token - client_id: os.environ/GITHUB_OAUTH_CLIENT_ID - client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET - scopes: ["public_repo", "user:email"] - allowed_tools: ["list_tools"] - # disallowed_tools: ["repo_delete"] + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY +litellm_settings: + callbacks: ["prometheus"] + custom_prometheus_metadata_labels: ["metadata.initiative"] \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c5370eb7d70..00ffae718e1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3066,6 +3066,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium = [ "tags", "team_member_key_duration", "prompts", + "logging", ] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e077d0ee923..73052d14957 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -579,7 +579,12 @@ class LiteLLMProxyRequestSetup: user_api_key_end_user_id=user_api_key_dict.end_user_id, user_api_key_user_email=user_api_key_dict.user_email, user_api_key_request_route=user_api_key_dict.request_route, - user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at.isoformat() if user_api_key_dict.budget_reset_at else None, + user_api_key_budget_reset_at=( + user_api_key_dict.budget_reset_at.isoformat() + if user_api_key_dict.budget_reset_at + else None + ), + user_api_key_auth_metadata=None, ) return user_api_key_logged_metadata @@ -607,6 +612,35 @@ class LiteLLMProxyRequestSetup: ) return data + @staticmethod + def add_management_endpoint_metadata_to_request_metadata( + data: dict, + management_endpoint_metadata: dict, + _metadata_variable_name: str, + ) -> dict: + """ + Adds the `UserAPIKeyAuth` metadata to the request metadata. + + ignore any sensitive fields like logging, api_key, etc. + """ + from litellm.proxy._types import ( + LiteLLM_ManagementEndpoint_MetadataFields, + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) + + # ignore any special fields + added_metadata = {} + for k, v in management_endpoint_metadata.items(): + if k not in ( + LiteLLM_ManagementEndpoint_MetadataFields_Premium + + LiteLLM_ManagementEndpoint_MetadataFields + ): + added_metadata[k] = v + data[_metadata_variable_name].setdefault( + "user_api_key_auth_metadata", {} + ).update(added_metadata) + return data + @staticmethod def add_key_level_controls( key_metadata: Optional[dict], data: dict, _metadata_variable_name: str @@ -651,6 +685,13 @@ class LiteLLMProxyRequestSetup: key_metadata["disable_fallbacks"], bool ): data["disable_fallbacks"] = key_metadata["disable_fallbacks"] + + ## KEY-LEVEL METADATA + data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=key_metadata, + _metadata_variable_name=_metadata_variable_name, + ) return data @staticmethod @@ -889,6 +930,15 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "spend_logs_metadata" ] + ## TEAM-LEVEL METADATA + data = ( + LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( + data=data, + management_endpoint_metadata=team_metadata, + _metadata_variable_name=_metadata_variable_name, + ) + ) + # Team spend, budget - used by prometheus.py data[_metadata_variable_name][ "user_api_key_team_max_budget" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e5786e50a5d..c8de97bba20 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -123,12 +123,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): max_output_tokens: Required[Optional[int]] input_cost_per_token: Required[float] input_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - input_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + input_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing cache_creation_input_token_cost: Optional[float] cache_creation_input_token_cost_above_1hr: Optional[float] cache_read_input_token_cost: Optional[float] - cache_read_input_token_cost_flex: Optional[float] # OpenAI flex service tier pricing - cache_read_input_token_cost_priority: Optional[float] # OpenAI priority service tier pricing + cache_read_input_token_cost_flex: Optional[ + float + ] # OpenAI flex service tier pricing + cache_read_input_token_cost_priority: Optional[ + float + ] # OpenAI priority service tier pricing input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models @@ -147,7 +153,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_batches: Optional[float] output_cost_per_token: Required[float] output_cost_per_token_flex: Optional[float] # OpenAI flex service tier pricing - output_cost_per_token_priority: Optional[float] # OpenAI priority service tier pricing + output_cost_per_token_priority: Optional[ + float + ] # OpenAI priority service tier pricing output_cost_per_character: Optional[float] # only for vertex ai models output_cost_per_audio_token: Optional[float] output_cost_per_token_above_128k_tokens: Optional[ @@ -1856,6 +1864,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_team_alias: Optional[str] user_api_key_end_user_id: Optional[str] user_api_key_request_route: Optional[str] + user_api_key_auth_metadata: Optional[Dict[str, str]] class StandardLoggingMCPToolCall(TypedDict, total=False): @@ -2059,10 +2068,12 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): StandardLoggingPayloadStatus = Literal["success", "failure"] + class CachingDetails(TypedDict): """ Track all caching related metrics, fields for a given request """ + cache_hit: Optional[bool] """ Whether the request hit the cache @@ -2072,12 +2083,16 @@ class CachingDetails(TypedDict): Duration for reading from cache """ + class CostBreakdown(TypedDict): """ Detailed cost breakdown for a request """ + input_cost: float # Cost of input/prompt tokens - output_cost: float # Cost of output/completion tokens (includes reasoning if applicable) + output_cost: ( + float # Cost of output/completion tokens (includes reasoning if applicable) + ) total_cost: float # Total cost (input + output + tool usage) tool_usage_cost: float # Cost of usage of built-in tools @@ -2616,6 +2631,7 @@ class SpecialEnums(Enum): class ServiceTier(Enum): """Enum for service tier types used in cost calculations.""" + FLEX = "flex" PRIORITY = "priority" @@ -2662,13 +2678,14 @@ CostResponseTypes = Union[ class PriorityReservationSettings(BaseModel): """ Settings for priority-based rate limiting reservation. - + Defines what priority to assign to keys without explicit priority metadata. The priority_reservation mapping is configured separately via litellm.priority_reservation. """ + default_priority: float = Field( default=0.5, - description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation." + description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation.", ) model_config = ConfigDict(protected_namespaces=()) From a1a0e99638ebca998db597649b679a4f1d869a81 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Sep 2025 21:23:25 -0700 Subject: [PATCH 018/474] fix(prometheus.py): working e2e calls w/ userapikeymetadata --- .../litellm_enterprise/integrations/prometheus.py | 11 +++++------ litellm/proxy/_new_secret_config.yaml | 2 +- litellm/proxy/litellm_pre_call_utils.py | 8 +++++--- litellm/types/integrations/prometheus.py | 8 ++++---- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index b472ccbb357..3b37e14b896 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -21,6 +21,7 @@ from litellm._logging import print_verbose, verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth from litellm.types.integrations.prometheus import * +from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_end_user_id_for_cost_tracking @@ -2247,12 +2248,10 @@ def prometheus_label_factory( if enum_values.custom_metadata_labels is not None: for key, value in enum_values.custom_metadata_labels.items(): - if key in supported_enum_labels: - filtered_labels[key] = value - else: - filtered_labels[key] = ( - "None" # this happens for dynamically added metadata labels - ) + # check sanitized key + sanitized_key = _sanitize_prometheus_label_name(key) + if sanitized_key in supported_enum_labels: + filtered_labels[sanitized_key] = value # Add custom tags if configured if enum_values.tags is not None: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 96d9ea0cc31..5d8052493a1 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -6,4 +6,4 @@ model_list: litellm_settings: callbacks: ["prometheus"] - custom_prometheus_metadata_labels: ["metadata.initiative"] \ No newline at end of file + custom_prometheus_metadata_labels: ["metadata.initiative", "metadata.business-unit"] \ No newline at end of file diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 73052d14957..44e26313f95 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -636,9 +636,11 @@ class LiteLLMProxyRequestSetup: + LiteLLM_ManagementEndpoint_MetadataFields ): added_metadata[k] = v - data[_metadata_variable_name].setdefault( - "user_api_key_auth_metadata", {} - ).update(added_metadata) + if data[_metadata_variable_name].get("user_api_key_auth_metadata") is None: + data[_metadata_variable_name]["user_api_key_auth_metadata"] = {} + data[_metadata_variable_name]["user_api_key_auth_metadata"].update( + added_metadata + ) return data @staticmethod diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 9c1a14a830e..a3dd4dcb1c6 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -426,13 +426,13 @@ class PrometheusMetricLabels: # Buffer monitoring metrics - these typically don't need additional labels litellm_pod_lock_manager_size: List[str] = [] - + litellm_in_memory_daily_spend_update_queue_size: List[str] = [] - + litellm_redis_daily_spend_update_queue_size: List[str] = [] - + litellm_in_memory_spend_update_queue_size: List[str] = [] - + litellm_redis_spend_update_queue_size: List[str] = [] @staticmethod From 933b3979eba92a14bacc84a682424c32fdcb1be3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Speglich?= Date: Wed, 1 Oct 2025 10:45:50 -0300 Subject: [PATCH 019/474] docs: update oci docs with oci_serving_mode --- docs/my-website/docs/providers/oci.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/my-website/docs/providers/oci.md b/docs/my-website/docs/providers/oci.md index 6fc1835154a..c11d64f4553 100644 --- a/docs/my-website/docs/providers/oci.md +++ b/docs/my-website/docs/providers/oci.md @@ -44,6 +44,7 @@ response = completion( oci_user=, oci_fingerprint=, oci_tenancy=, + oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED" # Provide either the private key string OR the path to the key file: # Option 1: pass the private key as a string oci_key=, @@ -71,6 +72,7 @@ response = completion( oci_user=, oci_fingerprint=, oci_tenancy=, + oci_serving_mode="ON_DEMAND", # Optional, default is "ON_DEMAND". Other option is "DEDICATED" # Provide either the private key string OR the path to the key file: # Option 1: pass the private key as a string oci_key=, From c32f42098cf94c2cfa44826a8769c4364993b658 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 19:57:16 +0530 Subject: [PATCH 020/474] Add cost tracking for /v1/messages --- litellm/proxy/common_request_processing.py | 83 +++++++++++- .../anthropic_passthrough_logging_handler.py | 5 +- .../test_anthropic_passthrough.py | 120 ++++++++++++++++++ 3 files changed, 205 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f07a61c544c..e7bad1e19dd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -734,7 +734,7 @@ class ProxyBaseLLMRequestProcessing: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events """ - from litellm.types.utils import ModelResponse, ModelResponseStream + from litellm.types.utils import ModelResponse, ModelResponseStream, Usage verbose_proxy_logger.debug("inside generator") try: @@ -759,6 +759,87 @@ class ProxyBaseLLMRequestProcessing: response_str = litellm.get_response_string(response_obj=chunk) str_so_far += response_str + # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider + # Handle both dict SSE events and pre-formatted string SSE lines + if getattr(litellm, "include_cost_in_streaming_usage", False) is True: + try: + def _inject_cost_into_usage_dict(obj: dict) -> Optional[dict]: + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + _mr = ModelResponse( + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + ) + model_name = request_data.get("model", "") + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None + + def _inject_cost_into_sse_frame_str(frame_str: str) -> Optional[str]: + # frame_str may contain multiple lines like 'event: ...\ndata: {...}\n\n' + # We only modify the JSON in the 'data:' line + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = _inject_cost_into_usage_dict(obj) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + if isinstance(chunk, dict): + maybe_modified = _inject_cost_into_usage_dict(chunk) + if maybe_modified is not None: + chunk = maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = _inject_cost_into_sse_frame_str(s) + if maybe_mod is not None: + chunk = (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = _inject_cost_into_sse_frame_str(chunk) + if maybe_mod is not None: + # Ensure trailing frame separator + chunk = maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index b9858202bf8..9b6e22b8196 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -20,7 +20,7 @@ from litellm.types.utils import ModelResponse, TextCompletionResponse if TYPE_CHECKING: from ..success_handler import PassThroughEndpointLogging - from ..types import EndpointType + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType else: PassThroughEndpointLogging = Any EndpointType = Any @@ -228,6 +228,7 @@ class AnthropicPassthroughLoggingHandler: except (StopIteration, StopAsyncIteration): break complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, ) return complete_streaming_response diff --git a/tests/pass_through_tests/test_anthropic_passthrough.py b/tests/pass_through_tests/test_anthropic_passthrough.py index 002fb20e9e8..6e819f9971c 100644 --- a/tests/pass_through_tests/test_anthropic_passthrough.py +++ b/tests/pass_through_tests/test_anthropic_passthrough.py @@ -296,3 +296,123 @@ async def test_anthropic_streaming_with_headers(): assert log_entry["end_user"] == "test-user-1" assert log_entry["custom_llm_provider"] == "anthropic" + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for Anthropic Messages API streaming + """ + print("Testing cost injection in Anthropic Messages API streaming response") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "claude-3-7-sonnet-20250219", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") + + +@pytest.mark.asyncio +@pytest.mark.flaky(retries=3, delay=2) +async def test_anthropic_messages_openai_model_streaming_cost_injection(): + """ + Test that cost is injected into message_delta usage for OpenAI model via Anthropic Messages API + """ + print("Testing cost injection in Anthropic Messages API with OpenAI model") + + headers = { + "Authorization": "Bearer sk-1234", + "Content-Type": "application/json", + "anthropic-version": "2023-06-01", + } + + payload = { + "model": "openai/gpt-4o", + "max_tokens": 10, + "stream": True, + "messages": [{"role": "user", "content": "Say 'Hi'"}], + } + + async with aiohttp.ClientSession() as session: + async with session.post( + "http://0.0.0.0:4000/v1/messages", + json=payload, + headers=headers + ) as response: + assert response.status == 200 + + # Collect all SSE events + events = [] + async for line in response.content: + line_str = line.decode('utf-8').strip() + if line_str.startswith('data: '): + try: + data = json.loads(line_str[6:]) # Remove 'data: ' prefix + events.append(data) + except json.JSONDecodeError: + continue + + # Find message_delta event with usage + message_delta_events = [ + event for event in events + if event.get('type') == 'message_delta' and 'usage' in event + ] + + assert len(message_delta_events) > 0, "No message_delta events with usage found" + + # Check that cost is included in usage + for event in message_delta_events: + usage = event.get('usage', {}) + assert 'cost' in usage, f"Cost not found in usage: {usage}" + assert isinstance(usage['cost'], (int, float)), f"Cost should be numeric: {usage['cost']}" + assert usage['cost'] >= 0, f"Cost should be non-negative: {usage['cost']}" + + print(f"✅ Found message_delta with cost: {usage}") + + print(f"✅ Test passed: Found {len(message_delta_events)} message_delta events with cost") From 56e429e33d036548e75ce4ded5e9a782c97b9684 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 1 Oct 2025 23:38:53 +0530 Subject: [PATCH 021/474] refactor code for better handling cost --- litellm/proxy/common_request_processing.py | 198 +++++++++++++-------- 1 file changed, 119 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e7bad1e19dd..091077b6e7e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -46,6 +46,7 @@ if TYPE_CHECKING: else: ProxyConfig = Any from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: @@ -760,85 +761,8 @@ class ProxyBaseLLMRequestProcessing: str_so_far += response_str # Inject cost into Anthropic-style SSE usage for /v1/messages for any provider - # Handle both dict SSE events and pre-formatted string SSE lines - if getattr(litellm, "include_cost_in_streaming_usage", False) is True: - try: - def _inject_cost_into_usage_dict(obj: dict) -> Optional[dict]: - if ( - obj.get("type") == "message_delta" - and isinstance(obj.get("usage"), dict) - ): - _usage = obj["usage"] - prompt_tokens = int(_usage.get("input_tokens", 0) or 0) - completion_tokens = int(_usage.get("output_tokens", 0) or 0) - total_tokens = int( - _usage.get("total_tokens", prompt_tokens + completion_tokens) - or (prompt_tokens + completion_tokens) - ) - - _mr = ModelResponse( - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=total_tokens, - ) - ) - model_name = request_data.get("model", "") - try: - cost_val = litellm.completion_cost( - completion_response=_mr, - model=model_name, - ) - except Exception: - cost_val = None - if cost_val is not None: - obj.setdefault("usage", {})["cost"] = cost_val - return obj - return None - - def _inject_cost_into_sse_frame_str(frame_str: str) -> Optional[str]: - # frame_str may contain multiple lines like 'event: ...\ndata: {...}\n\n' - # We only modify the JSON in the 'data:' line - try: - # Split preserving lines - lines = frame_str.split("\n") - for idx, ln in enumerate(lines): - stripped_ln = ln.strip() - if stripped_ln.startswith("data:"): - json_part = stripped_ln.split("data:", 1)[1].strip() - if json_part and json_part != "[DONE]": - obj = json.loads(json_part) - maybe_modified = _inject_cost_into_usage_dict(obj) - if maybe_modified is not None: - # Replace just this line with updated JSON using safe_dumps - lines[idx] = f"data: {safe_dumps(maybe_modified)}" - return "\n".join(lines) - return None - except Exception: - return None - - if isinstance(chunk, dict): - maybe_modified = _inject_cost_into_usage_dict(chunk) - if maybe_modified is not None: - chunk = maybe_modified - elif isinstance(chunk, (bytes, bytearray)): - # Decode to str, inject, and rebuild as bytes - try: - s = chunk.decode("utf-8", errors="ignore") - maybe_mod = _inject_cost_into_sse_frame_str(s) - if maybe_mod is not None: - chunk = (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") - except Exception: - pass - elif isinstance(chunk, str): - # Try to parse SSE frame and inject cost into the data line - maybe_mod = _inject_cost_into_sse_frame_str(chunk) - if maybe_mod is not None: - # Ensure trailing frame separator - chunk = maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") - except Exception: - # Never break streaming on optional cost injection - pass + model_name = request_data.get("model", "") + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(chunk, model_name) # Format chunk using helper function yield ProxyBaseLLMRequestProcessing.return_sse_chunk(chunk) @@ -871,3 +795,119 @@ class ProxyBaseLLMRequestProcessing: ) error_returned = json.dumps({"error": proxy_exception.to_dict()}) yield f"{STREAM_SSE_DATA_PREFIX}{error_returned}\n\n" + + @staticmethod + def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + """ + Process a streaming chunk and inject cost information if enabled. + + Args: + chunk: The streaming chunk (dict, str, bytes, or bytearray) + model_name: Model name for cost calculation + + Returns: + The processed chunk with cost information injected if applicable + """ + if not getattr(litellm, "include_cost_in_streaming_usage", False): + return chunk + + try: + if isinstance(chunk, dict): + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(chunk, model_name) + if maybe_modified is not None: + return maybe_modified + elif isinstance(chunk, (bytes, bytearray)): + # Decode to str, inject, and rebuild as bytes + try: + s = chunk.decode("utf-8", errors="ignore") + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name) + if maybe_mod is not None: + return (maybe_mod + ("" if maybe_mod.endswith("\n\n") else "\n\n")).encode("utf-8") + except Exception: + pass + elif isinstance(chunk, str): + # Try to parse SSE frame and inject cost into the data line + maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(chunk, model_name) + if maybe_mod is not None: + # Ensure trailing frame separator + return maybe_mod if maybe_mod.endswith("\n\n") else (maybe_mod + "\n\n") + except Exception: + # Never break streaming on optional cost injection + pass + + return chunk + + @staticmethod + def _inject_cost_into_sse_frame_str(frame_str: str, model_name: str) -> Optional[str]: + """ + Inject cost information into an SSE frame string by modifying the JSON in the 'data:' line. + + Args: + frame_str: SSE frame string that may contain multiple lines + model_name: Model name for cost calculation + + Returns: + Modified SSE frame string with cost injected, or None if no modification needed + """ + try: + # Split preserving lines + lines = frame_str.split("\n") + for idx, ln in enumerate(lines): + stripped_ln = ln.strip() + if stripped_ln.startswith("data:"): + json_part = stripped_ln.split("data:", 1)[1].strip() + if json_part and json_part != "[DONE]": + obj = json.loads(json_part) + maybe_modified = ProxyBaseLLMRequestProcessing._inject_cost_into_usage_dict(obj, model_name) + if maybe_modified is not None: + # Replace just this line with updated JSON using safe_dumps + lines[idx] = f"data: {safe_dumps(maybe_modified)}" + return "\n".join(lines) + return None + except Exception: + return None + + @staticmethod + def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]: + """ + Inject cost information into a usage dictionary for message_delta events. + + Args: + obj: Dictionary containing the SSE event data + model_name: Model name for cost calculation + + Returns: + Modified dictionary with cost injected, or None if no modification needed + """ + if ( + obj.get("type") == "message_delta" + and isinstance(obj.get("usage"), dict) + ): + _usage = obj["usage"] + prompt_tokens = int(_usage.get("input_tokens", 0) or 0) + completion_tokens = int(_usage.get("output_tokens", 0) or 0) + total_tokens = int( + _usage.get("total_tokens", prompt_tokens + completion_tokens) + or (prompt_tokens + completion_tokens) + ) + + _mr = ModelResponse( + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + ) + + try: + cost_val = litellm.completion_cost( + completion_response=_mr, + model=model_name, + ) + except Exception: + cost_val = None + + if cost_val is not None: + obj.setdefault("usage", {})["cost"] = cost_val + return obj + return None \ No newline at end of file From 005aec69c7879734b97f8f0c1c6d42c287528b22 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 16:57:07 -0700 Subject: [PATCH 022/474] feat(key_management_endpoints.py): allow specifying rate limit type when creating tpm/rpm limits on keys prevents overallocating tpm/rpm limits --- litellm/proxy/_types.py | 8 ++ .../key_management_endpoints.py | 100 +++++++++++++++--- 2 files changed, 91 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c5370eb7d70..d39f12f05e2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -731,6 +731,12 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): metadata: Optional[dict] = {} tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + rpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm + tpm_limit_type: Optional[ + Literal["guaranteed_throughput", "best_effort_throughput"] + ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm budget_duration: Optional[str] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} @@ -3054,6 +3060,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields = [ "model_rpm_limit", "model_tpm_limit", + "rpm_limit_type", + "tpm_limit_type", "guardrails", "tags", "enforced_params", diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 007c0164be4..3b39f609e1a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -90,10 +90,10 @@ def _get_user_in_team( def _calculate_key_rotation_time(rotation_interval: str) -> datetime: """ Helper function to calculate the next rotation time for a key based on the rotation interval. - + Args: rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h') - + Returns: datetime: The calculated next rotation time in UTC """ @@ -102,21 +102,25 @@ def _calculate_key_rotation_time(rotation_interval: str) -> datetime: return now + timedelta(seconds=interval_seconds) -def _set_key_rotation_fields(data: dict, auto_rotate: bool, rotation_interval: Optional[str]) -> None: +def _set_key_rotation_fields( + data: dict, auto_rotate: bool, rotation_interval: Optional[str] +) -> None: """ Helper function to set rotation fields in key data if auto_rotate is enabled. - + Args: data: Dictionary to update with rotation fields auto_rotate: Whether auto rotation is enabled rotation_interval: The rotation interval string (required if auto_rotate is True) """ if auto_rotate and rotation_interval: - data.update({ - "auto_rotate": auto_rotate, - "rotation_interval": rotation_interval, - "key_rotation_at": _calculate_key_rotation_time(rotation_interval) - }) + data.update( + { + "auto_rotate": auto_rotate, + "rotation_interval": rotation_interval, + "key_rotation_at": _calculate_key_rotation_time(rotation_interval), + } + ) def _is_allowed_to_make_key_request( @@ -542,6 +546,15 @@ async def _common_key_generation_helper( # noqa: PLR0915 value=getattr(data, field), ) + for field in LiteLLM_ManagementEndpoint_MetadataFields: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=data, + field_name=field, + value=getattr(data, field), + ) + delattr(data, field) + data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore data_json = handle_key_type(data, data_json) @@ -620,6 +633,46 @@ async def _common_key_generation_helper( # noqa: PLR0915 return response +async def _check_team_key_limits( + team_table: LiteLLM_TeamTableCachedObj, + data: GenerateKeyRequest, + prisma_client: PrismaClient, +) -> None: + """ + Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + """ + # get all team keys + # calculate allocated tpm/rpm limit + # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": team_table.team_id}, + ) + if keys is not None and len(keys) > 0: + allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None) + allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None) + else: + allocated_tpm = 0 + allocated_rpm = 0 + if ( + data.tpm_limit is not None + and team_table.tpm_limit is not None + and data.tpm_limit + allocated_tpm > team_table.tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={allocated_tpm} + Key TPM limit={data.tpm_limit} is greater than team TPM limit={team_table.tpm_limit}", + ) + if ( + data.rpm_limit is not None + and team_table.rpm_limit is not None + and data.rpm_limit + allocated_rpm > team_table.rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={allocated_rpm} + Key RPM limit={data.rpm_limit} is greater than team RPM limit={team_table.rpm_limit}", + ) + + @router.post( "/key/generate", tags=["key management"], @@ -696,12 +749,19 @@ async def generate_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ try: + from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, user_custom_key_generate, ) + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + verbose_proxy_logger.debug("entered /key/generate") if user_custom_key_generate is not None: @@ -729,7 +789,6 @@ async def generate_key_fn( verbose_proxy_logger.debug( f"Error getting team object in `/key/generate`: {e}" ) - team_table = None key_generation_check( team_table=team_table, @@ -738,12 +797,21 @@ async def generate_key_fn( route=KeyManagementRoutes.KEY_GENERATE, ) + if team_table is not None: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=prisma_client, + ) + return await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, team_table=team_table, ) + except HTTPException as e: + raise e except Exception as e: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.generate_key_fn(): Exception occured - {}".format( @@ -1198,9 +1266,9 @@ async def update_key_fn( # Handle rotation fields if auto_rotate is being enabled _set_key_rotation_fields( - non_default_values, - non_default_values.get("auto_rotate", False), - non_default_values.get("rotation_interval") + non_default_values, + non_default_values.get("auto_rotate", False), + non_default_values.get("rotation_interval"), ) _data = {**non_default_values, "token": key} @@ -1602,8 +1670,6 @@ def _check_model_access_group( return True - - async def generate_key_helper_fn( # noqa: PLR0915 request_type: Literal[ "user", "key" @@ -1766,12 +1832,12 @@ async def generate_key_helper_fn( # noqa: PLR0915 "allowed_routes": allowed_routes or [], "object_permission_id": object_permission_id, } - + # Add rotation fields if auto_rotate is enabled _set_key_rotation_fields( data=key_data, auto_rotate=auto_rotate or False, - rotation_interval=rotation_interval + rotation_interval=rotation_interval, ) if ( From a83238a2db788c81637581f168b5ff88bbb3b834 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:02:05 -0700 Subject: [PATCH 023/474] test: add unit tests --- .../test_key_management_endpoints.py | 520 +++++++++++++++--- 1 file changed, 448 insertions(+), 72 deletions(-) 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 e3aa7d58872..35ecdd8e0ae 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 @@ -15,12 +15,14 @@ from fastapi import HTTPException from litellm.proxy._types import ( GenerateKeyRequest, + LiteLLM_TeamTableCachedObj, LiteLLM_VerificationToken, LitellmUserRoles, UpdateKeyRequest, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_team_key_limits, _common_key_generation_helper, _list_key_helper, generate_key_helper_fn, @@ -1040,7 +1042,7 @@ async def test_unblock_key_invalid_key_format(monkeypatch): def test_validate_key_team_change_with_member_permissions(): """ Test validate_key_team_change function with team member permissions. - + This test covers the new logic that allows team members with specific permissions to update keys, not just team admins. """ @@ -1054,111 +1056,107 @@ def test_validate_key_team_change_with_member_permissions(): mock_key.models = ["gpt-4"] mock_key.tpm_limit = None mock_key.rpm_limit = None - + mock_team = MagicMock() - mock_team.team_id = "test-team-456" + mock_team.team_id = "test-team-456" mock_team.members_with_roles = [] mock_team.tpm_limit = None mock_team.rpm_limit = None - + mock_change_initiator = MagicMock() mock_change_initiator.user_id = "test-user-123" - + mock_router = MagicMock() - + # Mock the member object returned by _get_user_in_team mock_member_object = MagicMock() - - with patch('litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model'): - with patch('litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team') as mock_get_user: - with patch('litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin') as mock_is_admin: - with patch('litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint') as mock_has_perms: - + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model" + ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin" + ) as mock_is_admin: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint" + ) as mock_has_perms: + mock_get_user.return_value = mock_member_object mock_is_admin.return_value = False mock_has_perms.return_value = True - + # This should not raise an exception due to member permissions validate_key_team_change( key=mock_key, team=mock_team, change_initiated_by=mock_change_initiator, - llm_router=mock_router + llm_router=mock_router, ) - + # Verify the permission check was called with correct parameters mock_has_perms.assert_called_once_with( team_member_object=mock_member_object, team_table=mock_team, - route=KeyManagementRoutes.KEY_UPDATE.value + route=KeyManagementRoutes.KEY_UPDATE.value, ) def test_key_rotation_fields_helper(): """ Test the key data update logic for rotation fields. - + This test focuses on the core logic that adds rotation fields to key_data when auto_rotate is enabled, without the complexity of full key generation. """ # Test Case 1: With rotation enabled - key_data = { - "models": ["gpt-3.5-turbo"], - "user_id": "test-user" - } - + key_data = {"models": ["gpt-3.5-turbo"], "user_id": "test-user"} + auto_rotate = True rotation_interval = "30d" - + # Simulate the rotation logic from generate_key_helper_fn if auto_rotate and rotation_interval: - key_data.update({ - "auto_rotate": auto_rotate, - "rotation_interval": rotation_interval - }) - + key_data.update( + {"auto_rotate": auto_rotate, "rotation_interval": rotation_interval} + ) + # Verify rotation fields are added assert key_data["auto_rotate"] == True assert key_data["rotation_interval"] == "30d" assert key_data["models"] == ["gpt-3.5-turbo"] # Original fields preserved - + # Test Case 2: Without rotation enabled - key_data2 = { - "models": ["gpt-4"], - "user_id": "test-user" - } - + key_data2 = {"models": ["gpt-4"], "user_id": "test-user"} + auto_rotate2 = False rotation_interval2 = None - + # Simulate the rotation logic if auto_rotate2 and rotation_interval2: - key_data2.update({ - "auto_rotate": auto_rotate2, - "rotation_interval": rotation_interval2 - }) - + key_data2.update( + {"auto_rotate": auto_rotate2, "rotation_interval": rotation_interval2} + ) + # Verify rotation fields are NOT added assert "auto_rotate" not in key_data2 assert "rotation_interval" not in key_data2 assert key_data2["models"] == ["gpt-4"] # Original fields preserved - + # Test Case 3: auto_rotate=True but no interval - key_data3 = { - "models": ["claude-3"], - "user_id": "test-user" - } - + key_data3 = {"models": ["claude-3"], "user_id": "test-user"} + auto_rotate3 = True rotation_interval3 = None - + # Simulate the rotation logic if auto_rotate3 and rotation_interval3: - key_data3.update({ - "auto_rotate": auto_rotate3, - "rotation_interval": rotation_interval3 - }) - + key_data3.update( + {"auto_rotate": auto_rotate3, "rotation_interval": rotation_interval3} + ) + # Verify rotation fields are NOT added (missing interval) assert "auto_rotate" not in key_data3 assert "rotation_interval" not in key_data3 @@ -1181,27 +1179,24 @@ async def test_update_key_fn_auto_rotate_enable(): team_id=None, auto_rotate=False, rotation_interval=None, - metadata={} + metadata={}, ) - + # Test enabling auto rotation update_request = UpdateKeyRequest( - key="test-token", - auto_rotate=True, - rotation_interval="30d" + key="test-token", auto_rotate=True, rotation_interval="30d" ) - + result = await prepare_key_update_data( - data=update_request, - existing_key_row=existing_key + data=update_request, existing_key_row=existing_key ) - + # Verify rotation fields are included assert result["auto_rotate"] is True assert result["rotation_interval"] == "30d" -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_update_key_fn_auto_rotate_disable(): """Test that update_key_fn properly handles disabling auto rotation.""" from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest @@ -1218,19 +1213,400 @@ async def test_update_key_fn_auto_rotate_disable(): team_id=None, auto_rotate=True, rotation_interval="30d", - metadata={} + metadata={}, ) - + # Test disabling auto rotation - update_request = UpdateKeyRequest( - key="test-token", - auto_rotate=False - ) - + update_request = UpdateKeyRequest(key="test-token", auto_rotate=False) + result = await prepare_key_update_data( - data=update_request, - existing_key_row=existing_key + data=update_request, existing_key_row=existing_key ) - + # Verify auto_rotate is set to False assert result["auto_rotate"] is False + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_existing_keys(): + """ + Test _check_team_key_limits when team has no existing keys. + Should allow any TPM/RPM limits within team bounds. + """ + # Mock prisma client + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with limits within team bounds + data = GenerateKeyRequest( + tpm_limit=5000, + rpm_limit=500, + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + # Verify database was queried + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={"team_id": "test-team-123"} + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_with_existing_keys_within_bounds(): + """ + Test _check_team_key_limits when team has existing keys but total allocation + is still within team limits. + """ + # Create mock existing keys + existing_key1 = MagicMock() + existing_key1.tpm_limit = 3000 + existing_key1.rpm_limit = 200 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 2000 + existing_key2.rpm_limit = 300 + + existing_key3 = MagicMock() + existing_key3.tpm_limit = None # Should be ignored in calculation + existing_key3.rpm_limit = None # Should be ignored in calculation + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2, existing_key3] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-456", + team_alias="test-team", + tpm_limit=10000, # Total: 3000 + 2000 + 4000 (new) = 9000 < 10000 ✓ + rpm_limit=1000, # Total: 200 + 300 + 400 (new) = 900 < 1000 ✓ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would still be within bounds + data = GenerateKeyRequest( + tpm_limit=4000, + rpm_limit=400, + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_tpm_overallocation(): + """ + Test _check_team_key_limits when new key would cause TPM overallocation. + Should raise HTTPException with appropriate error message. + """ + # Create mock existing keys with high TPM usage + existing_key1 = MagicMock() + existing_key1.tpm_limit = 6000 + existing_key1.rpm_limit = 100 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 3000 + existing_key2.rpm_limit = 200 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-789", + team_alias="test-team", + tpm_limit=10000, # Allocated: 6000 + 3000 = 9000, New: 2000, Total: 11000 > 10000 ✗ + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would exceed TPM limits + data = GenerateKeyRequest( + tpm_limit=2000, + rpm_limit=100, + ) + + # Should raise HTTPException for TPM overallocation + with pytest.raises(HTTPException) as exc_info: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated TPM limit=9000 + Key TPM limit=2000 is greater than team TPM limit=10000" + in str(exc_info.value.detail) + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_rpm_overallocation(): + """ + Test _check_team_key_limits when new key would cause RPM overallocation. + Should raise HTTPException with appropriate error message. + """ + # Create mock existing keys with high RPM usage + existing_key1 = MagicMock() + existing_key1.tpm_limit = 1000 + existing_key1.rpm_limit = 600 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 2000 + existing_key2.rpm_limit = 300 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-101", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, # Allocated: 600 + 300 = 900, New: 200, Total: 1100 > 1000 ✗ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would exceed RPM limits + data = GenerateKeyRequest( + tpm_limit=1000, + rpm_limit=200, + ) + + # Should raise HTTPException for RPM overallocation + with pytest.raises(HTTPException) as exc_info: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated RPM limit=900 + Key RPM limit=200 is greater than team RPM limit=1000" + in str(exc_info.value.detail) + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_team_limits(): + """ + Test _check_team_key_limits when team has no TPM/RPM limits set. + Should allow any key limits since there are no team constraints. + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 5000 + existing_key.rpm_limit = 500 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with no limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-202", + team_alias="test-team", + tpm_limit=None, # No team limit + rpm_limit=None, # No team limit + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with any limits + data = GenerateKeyRequest( + tpm_limit=10000, # High limit should be allowed + rpm_limit=2000, # High limit should be allowed + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_key_limits(): + """ + Test _check_team_key_limits when new key has no TPM/RPM limits. + Should not raise any exceptions since no limits are being allocated. + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 8000 + existing_key.rpm_limit = 800 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-303", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with no limits + data = GenerateKeyRequest( + tpm_limit=None, # No limit being set + rpm_limit=None, # No limit being set + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_mixed_scenarios(): + """ + Test _check_team_key_limits with mixed scenarios: + - Some existing keys have limits, others don't + - New key has only one type of limit + - Team has only one type of limit + """ + # Create mock existing keys with mixed limits + existing_key1 = MagicMock() + existing_key1.tpm_limit = 3000 + existing_key1.rpm_limit = None # No RPM limit + + existing_key2 = MagicMock() + existing_key2.tpm_limit = None # No TPM limit + existing_key2.rpm_limit = 400 + + existing_key3 = MagicMock() + existing_key3.tpm_limit = 2000 + existing_key3.rpm_limit = 300 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2, existing_key3] + ) + + # Create team table with only TPM limit + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-404", + team_alias="test-team", + tpm_limit=10000, # Allocated: 3000 + 0 + 2000 = 5000, New: 4000, Total: 9000 < 10000 ✓ + rpm_limit=None, # No team RPM limit + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with only TPM limit + data = GenerateKeyRequest( + tpm_limit=4000, + rpm_limit=None, # No RPM limit being set + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_exact_boundary(): + """ + Test _check_team_key_limits when allocation exactly matches team limits. + Should allow the allocation (boundary case). + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 7000 + existing_key.rpm_limit = 700 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-505", + team_alias="test-team", + tpm_limit=10000, # Allocated: 7000, New: 3000, Total: 10000 = 10000 ✓ + rpm_limit=1000, # Allocated: 700, New: 300, Total: 1000 = 1000 ✓ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that exactly matches remaining capacity + data = GenerateKeyRequest( + tpm_limit=3000, + rpm_limit=300, + ) + + # Should not raise any exception (exact boundary should be allowed) + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) From 3ce074564e6202b1dd65710f6d8f56c7662aa932 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:48:23 -0700 Subject: [PATCH 024/474] feat(key_management_endpoints.py): add guaranteed throughput support for model specific tpm / rpm limits prevents admin from granting keys more tpm/rpm than created for a key --- .../key_management_endpoints.py | 118 ++++++++++++++++-- 1 file changed, 109 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3b39f609e1a..b0f9c1573d5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -27,6 +27,7 @@ from litellm.caching import DualCache from litellm.constants import LENGTH_OF_LITELLM_GENERATED_KEY, UI_SESSION_TOKEN_TEAM_ID from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * +from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.auth.auth_checks import ( _cache_key_object, _delete_cache_key_object, @@ -633,20 +634,93 @@ async def _common_key_generation_helper( # noqa: PLR0915 return response -async def _check_team_key_limits( +def check_team_key_model_specific_limits( + keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, data: GenerateKeyRequest, - prisma_client: PrismaClient, ) -> None: """ - Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating. + """ + if data.model_rpm_limit is None and data.model_tpm_limit is None: + return + # get total model specific tpm/rpm limit + model_specific_rpm_limit = {} + model_specific_tpm_limit = {} + + for key in keys: + if key.metadata.get("model_rpm_limit", None) is not None: + for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items(): + model_specific_rpm_limit[model] = ( + model_specific_rpm_limit.get(model, 0) + rpm_limit + ) + if key.metadata.get("model_tpm_limit", None) is not None: + for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items(): + model_specific_tpm_limit[model] = ( + model_specific_tpm_limit.get(model, 0) + tpm_limit + ) + if data.model_rpm_limit is not None: + for model, rpm_limit in data.model_rpm_limit.items(): + if ( + model_specific_rpm_limit.get(model, 0) + rpm_limit + > team_table.rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than team RPM limit={team_table.rpm_limit}", + ) + elif team_table.metadata and team_table.metadata.get("model_rpm_limit"): + team_model_specific_rpm_limit_dict = team_table.metadata.get( + "model_rpm_limit", {} + ) + team_model_specific_rpm_limit = team_model_specific_rpm_limit_dict.get( + model + ) + if ( + model_specific_rpm_limit.get(model, 0) + rpm_limit + > team_model_specific_rpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than team RPM limit={team_model_specific_rpm_limit.get(model, 0)}", + ) + if data.model_tpm_limit is not None: + for model, tpm_limit in data.model_tpm_limit.items(): + if ( + team_table.tpm_limit is not None + and model_specific_tpm_limit.get(model, 0) + tpm_limit + > team_table.tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than team TPM limit={team_table.tpm_limit}", + ) + elif team_table.metadata and team_table.metadata.get("model_tpm_limit"): + team_model_specific_tpm_limit_dict = team_table.metadata.get( + "model_tpm_limit", {} + ) + team_model_specific_tpm_limit = team_model_specific_tpm_limit_dict.get( + model + ) + if ( + team_model_specific_tpm_limit + and model_specific_tpm_limit.get(model, 0) + tpm_limit + > team_model_specific_tpm_limit + ): + raise HTTPException( + status_code=400, + detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than team TPM limit={team_model_specific_tpm_limit}", + ) + + +def check_team_key_rpm_tpm_limits( + keys: List[LiteLLM_VerificationToken], + team_table: LiteLLM_TeamTableCachedObj, + data: GenerateKeyRequest, +) -> None: + """ + Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. """ - # get all team keys - # calculate allocated tpm/rpm limit - # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": team_table.team_id}, - ) if keys is not None and len(keys) > 0: allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None) allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None) @@ -673,6 +747,32 @@ async def _check_team_key_limits( ) +async def _check_team_key_limits( + team_table: LiteLLM_TeamTableCachedObj, + data: GenerateKeyRequest, + prisma_client: PrismaClient, +) -> None: + """ + Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + """ + # get all team keys + # calculate allocated tpm/rpm limit + # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit + keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": team_table.team_id}, + ) + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=data, + ) + check_team_key_rpm_tpm_limits( + keys=keys, + team_table=team_table, + data=data, + ) + + @router.post( "/key/generate", tags=["key management"], From 757436f30235423074be0563e331f78d5af0752c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:51:10 -0700 Subject: [PATCH 025/474] test: add unit test --- .../test_key_management_endpoints.py | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) 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 35ecdd8e0ae..0b659acd2bb 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 @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_team_key_limits, _common_key_generation_helper, _list_key_helper, + check_team_key_model_specific_limits, generate_key_helper_fn, prepare_key_update_data, validate_key_team_change, @@ -1610,3 +1611,115 @@ async def test_check_team_key_limits_exact_boundary(): data=data, prisma_client=mock_prisma_client, ) + + +def test_check_team_key_model_specific_limits_no_limits(): + """ + Test check_team_key_model_specific_limits when no model-specific limits are set. + Should return without raising any exceptions. + """ + # Create existing key with no model-specific limits + existing_key = LiteLLM_VerificationToken( + token="test-token-1", + user_id="test-user", + team_id="test-team-123", + metadata={}, + ) + + keys = [existing_key] + + # Create team table + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + metadata={}, + ) + + # Create request with no model-specific limits + data = GenerateKeyRequest( + model_rpm_limit=None, + model_tpm_limit=None, + ) + + # Should not raise any exception + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=data, + ) + + +def test_check_team_key_model_specific_limits_rpm_overallocation(): + """ + Test check_team_key_model_specific_limits when model-specific RPM would cause overallocation. + Should raise HTTPException with appropriate error message. + """ + # Create existing keys with model-specific RPM limits + existing_key1 = LiteLLM_VerificationToken( + token="test-token-1", + user_id="test-user-1", + team_id="test-team-456", + metadata={ + "model_rpm_limit": { + "gpt-4": 500, + "gpt-3.5-turbo": 300, + } + }, + ) + + existing_key2 = LiteLLM_VerificationToken( + token="test-token-2", + user_id="test-user-2", + team_id="test-team-456", + metadata={ + "model_rpm_limit": { + "gpt-4": 300, + } + }, + ) + + keys = [existing_key1, existing_key2] + + # Create team table with RPM limit + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-456", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, # Total team RPM limit + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + metadata={}, + ) + + # Create request that would exceed model-specific RPM limits + # Existing gpt-4: 500 + 300 = 800, New: 300, Total: 1100 > 1000 (team limit) + data = GenerateKeyRequest( + model_rpm_limit={ + "gpt-4": 300, # This would cause overallocation + }, + model_tpm_limit=None, + ) + + # Should raise HTTPException for model-specific RPM overallocation + with pytest.raises(HTTPException) as exc_info: + check_team_key_model_specific_limits( + keys=keys, + team_table=team_table, + data=data, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated RPM limit=800 + Key RPM limit=300 is greater than team RPM limit=1000" + in str(exc_info.value.detail) + ) From d8e3a62fbe29af91ad43fdc0d1550fb5e824f69c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 17:58:31 -0700 Subject: [PATCH 026/474] feat(key_management_endpoints.py): add support for guaranteed throughput on key update and service account key creation --- .../key_management_endpoints.py | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b0f9c1573d5..22f3f5c530d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -637,7 +637,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 def check_team_key_model_specific_limits( keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: """ Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating. @@ -716,7 +716,7 @@ def check_team_key_model_specific_limits( def check_team_key_rpm_tpm_limits( keys: List[LiteLLM_VerificationToken], team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + data: Union[GenerateKeyRequest, UpdateKeyRequest], ) -> None: """ Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating. @@ -749,15 +749,23 @@ def check_team_key_rpm_tpm_limits( async def _check_team_key_limits( team_table: LiteLLM_TeamTableCachedObj, - data: GenerateKeyRequest, + data: Union[GenerateKeyRequest, UpdateKeyRequest], prisma_client: PrismaClient, ) -> None: """ Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating. + + Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput" """ + if ( + data.tpm_limit_type != "guaranteed_throughput" + and data.rpm_limit_type != "guaranteed_throughput" + ): + return # get all team keys # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit + keys = await prisma_client.db.litellm_verificationtoken.find_many( where={"team_id": team_table.team_id}, ) @@ -993,12 +1001,19 @@ async def generate_service_account_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ + from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, user_custom_key_generate, ) + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + await validate_team_id_used_in_service_account_request( team_id=data.team_id, prisma_client=prisma_client, @@ -1031,6 +1046,13 @@ async def generate_service_account_key_fn( ) team_table = None + if team_table is not None: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=prisma_client, + ) + key_generation_check( team_table=team_table, user_api_key_dict=user_api_key_dict, @@ -1330,14 +1352,22 @@ async def update_key_fn( user_api_key_cache=user_api_key_cache, ) + team_obj = await get_team_object( + team_id=cast(str, data.team_id), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + ) + + if team_obj is not None: + await _check_team_key_limits( + team_table=team_obj, + data=data, + prisma_client=prisma_client, + ) + # if team change - check if this is possible if is_different_team(data=data, existing_key_row=existing_key_row): - team_obj = await get_team_object( - team_id=cast(str, data.team_id), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) if llm_router is None: raise HTTPException( status_code=400, From ff866aae7bf5964a61fc4ca5c81c9524ba83ea69 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 18:30:34 -0700 Subject: [PATCH 027/474] feat(create_key_button.tsx): add tpm/rpm rate limit type options to UI allows user to set the type of tpm/rpm limit they're trying to set --- .../organisms/create_key_button.tsx | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) 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 dc523da7667..0ab8b25c1fe 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -810,6 +810,46 @@ const CreateKey: React.FC = ({ > + + TPM Rate Limit Type {' '} + + + + + } + name="tpm_limit_type" + initialValue="default" + className="mt-4" + > + + = ({ > + + RPM Rate Limit Type {' '} + + + + + } + name="rpm_limit_type" + initialValue="default" + className="mt-4" + > + + From 20b6f011f7eeae2131c15667037a60e6a50b26e1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 1 Oct 2025 18:34:29 -0700 Subject: [PATCH 028/474] fix(create_key_button.tsx): working ui controls to set guaranteed throughput/best effort throughput on key --- .../src/components/organisms/create_key_button.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 0ab8b25c1fe..e367565bae0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -820,7 +820,7 @@ const CreateKey: React.FC = ({ } name="tpm_limit_type" - initialValue="default" + initialValue={null} className="mt-4" > = ({ form.setFieldValue('rpm_limit_type', value); }} > -