From ff70dfef0037fbbf7f2db0f69fced6b11e8c7762 Mon Sep 17 00:00:00 2001 From: Honam <116502099+binbor111@users.noreply.github.com> Date: Thu, 2 Apr 2026 06:16:10 -0700 Subject: [PATCH] style: apply ruff formatting to modified files --- .../get_llm_provider_logic.py | 349 ++++-------------- tests/local_testing/test_get_llm_provider.py | 119 +++--- 2 files changed, 133 insertions(+), 335 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index d9c5e35286d..062add481bc 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -11,10 +11,7 @@ from ..types.router import LiteLLM_Params def _is_non_openai_azure_model(model: str) -> bool: try: model_name = model.split("/", 1)[1] - if ( - model_name in litellm.cohere_chat_models - or f"mistral/{model_name}" in litellm.mistral_chat_models - ): + if model_name in litellm.cohere_chat_models or f"mistral/{model_name}" in litellm.mistral_chat_models: return True except Exception: return False @@ -53,11 +50,7 @@ def handle_cohere_chat_model_custom_llm_provider( if model and "/" in model: _custom_llm_provider, _model = model.split("/", 1) - if ( - _custom_llm_provider - and _custom_llm_provider == "cohere" - and _model in litellm.cohere_chat_models - ): + if _custom_llm_provider and _custom_llm_provider == "cohere" and _model in litellm.cohere_chat_models: return _model, "cohere_chat" return model, custom_llm_provider @@ -78,10 +71,7 @@ def handle_anthropic_text_model_custom_llm_provider( """ if custom_llm_provider: - if ( - custom_llm_provider == "anthropic" - and litellm.AnthropicTextConfig._is_anthropic_text_model(model) - ): + if custom_llm_provider == "anthropic" and litellm.AnthropicTextConfig._is_anthropic_text_model(model): return model, "anthropic_text" if model and "/" in model: @@ -115,13 +105,9 @@ def get_llm_provider( # noqa: PLR0915 try: # Early validation - model is required if model is None: - raise ValueError( - "model parameter is required but was None. Please provide a valid model name." - ) + raise ValueError("model parameter is required but was None. Please provide a valid model name.") - if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default( - litellm_params=litellm_params - ): + if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default(litellm_params=litellm_params): return litellm.LiteLLMProxyChatConfig.litellm_proxy_get_custom_llm_provider_info( model=model, api_base=api_base, api_key=api_key ) @@ -145,13 +131,9 @@ def get_llm_provider( # noqa: PLR0915 return model, custom_llm_provider, dynamic_api_key, api_base ### Handle cases when custom_llm_provider is set to cohere/command-r-plus but it should use cohere_chat route - model, custom_llm_provider = handle_cohere_chat_model_custom_llm_provider( - model, custom_llm_provider - ) + model, custom_llm_provider = handle_cohere_chat_model_custom_llm_provider(model, custom_llm_provider) - model, custom_llm_provider = handle_anthropic_text_model_custom_llm_provider( - model, custom_llm_provider - ) + model, custom_llm_provider = handle_anthropic_text_model_custom_llm_provider(model, custom_llm_provider) if custom_llm_provider and ( model.split("/")[0] != custom_llm_provider @@ -197,9 +179,7 @@ def get_llm_provider( # noqa: PLR0915 custom_llm_provider = model.split("/", 1)[0] model = model.split("/", 1)[1] if api_base is not None and not isinstance(api_base, str): - raise Exception( - "api base needs to be a string. api_base={}".format(api_base) - ) + raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception( "dynamic_api_key needs to be a string. Got type={}".format( @@ -258,9 +238,7 @@ def get_llm_provider( # noqa: PLR0915 dynamic_api_key = get_secret_str("OLLAMA_API_KEY") elif endpoint == "https://api.friendli.ai/serverless/v1": custom_llm_provider = "friendliai" - dynamic_api_key = get_secret_str( - "FRIENDLIAI_API_KEY" - ) or get_secret("FRIENDLI_TOKEN") + dynamic_api_key = get_secret_str("FRIENDLIAI_API_KEY") or get_secret("FRIENDLI_TOKEN") elif endpoint == "api.galadriel.com/v1": custom_llm_provider = "galadriel" dynamic_api_key = get_secret_str("GALADRIEL_API_KEY") @@ -279,16 +257,10 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "api.moonshot.ai/v1": custom_llm_provider = "moonshot" dynamic_api_key = get_secret_str("MOONSHOT_API_KEY") - elif ( - endpoint == "api.minimax.io/anthropic" - or endpoint == "api.minimaxi.com/anthropic" - ): + elif endpoint == "api.minimax.io/anthropic" or endpoint == "api.minimaxi.com/anthropic": custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") - elif ( - endpoint == "api.minimax.io/v1" - or endpoint == "api.minimaxi.com/v1" - ): + elif endpoint == "api.minimax.io/v1" or endpoint == "api.minimaxi.com/v1": custom_llm_provider = "minimax" dynamic_api_key = get_secret_str("MINIMAX_API_KEY") elif endpoint == "platform.publicai.co/v1": @@ -329,18 +301,10 @@ def get_llm_provider( # noqa: PLR0915 dynamic_api_key = get_secret_str("WANDB_API_KEY") if api_base is not None and not isinstance(api_base, str): + raise Exception("api base needs to be a string. api_base={}".format(api_base)) + if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception( - "api base needs to be a string. api_base={}".format( - api_base - ) - ) - if dynamic_api_key is not None and not isinstance( - dynamic_api_key, str - ): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) + "dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key) ) return model, custom_llm_provider, dynamic_api_key, api_base # type: ignore @@ -369,13 +333,10 @@ def get_llm_provider( # noqa: PLR0915 elif model in litellm.cohere_chat_models: custom_llm_provider = "cohere_chat" ## replicate - elif model in litellm.replicate_models or ( - ":" in model and len(model) > REPLICATE_MODEL_NAME_WITH_ID_LENGTH - ): + elif model in litellm.replicate_models or (":" in model and len(model) > REPLICATE_MODEL_NAME_WITH_ID_LENGTH): model_parts = model.split(":") if ( - len(model_parts) > 1 - and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH + len(model_parts) > 1 and len(model_parts[1]) == REPLICATE_MODEL_NAME_WITH_ID_LENGTH ): ## checks if model name has a 64 digit code - e.g. "meta/llama-2-70b-chat:02e509c789964a7ea8736978a43525956ef40397be9033abf9fd2badfe68c9e3" custom_llm_provider = "replicate" elif model in litellm.replicate_models: @@ -402,11 +363,7 @@ def get_llm_provider( # noqa: PLR0915 ## ai21 elif model in litellm.ai21_chat_models or model in litellm.ai21_models: custom_llm_provider = "ai21_chat" - api_base = ( - api_base - or get_secret("AI21_API_BASE") - or "https://api.ai21.com/studio/v1" - ) # type: ignore + api_base = api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" # type: ignore dynamic_api_key = api_key or get_secret("AI21_API_KEY") ## aleph_alpha elif model in litellm.aleph_alpha_models: @@ -478,23 +435,15 @@ def get_llm_provider( # noqa: PLR0915 llm_provider="", ) if api_base is not None and not isinstance(api_base, str): - raise Exception( - "api base needs to be a string. api_base={}".format(api_base) - ) + raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) - ) + raise Exception("dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key)) return model, custom_llm_provider, dynamic_api_key, api_base except Exception as e: if isinstance(e, litellm.exceptions.BadRequestError): raise e else: - error_str = ( - f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}" - ) + error_str = f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}" raise litellm.exceptions.BadRequestError( # type: ignore message=f"GetLLMProvider Exception - {str(e)}\n\noriginal model: {model}", model=model, @@ -530,9 +479,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 if provider_config is None: raise ValueError(f"Provider {custom_llm_provider} not found") config_class = create_config_class(provider_config) - api_base, dynamic_api_key = config_class()._get_openai_compatible_provider_info( - api_base, api_key - ) + api_base, dynamic_api_key = config_class()._get_openai_compatible_provider_info(api_base, api_key) return model, custom_llm_provider, dynamic_api_key, api_base if custom_llm_provider == "perplexity": @@ -540,9 +487,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ( api_base, dynamic_api_key, - ) = litellm.PerplexityChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.PerplexityChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "aiohttp_openai": return model, "aiohttp_openai", api_key, api_base elif custom_llm_provider == "anyscale": @@ -553,152 +498,90 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ( api_base, dynamic_api_key, - ) = litellm.DeepInfraConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DeepInfraConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "empower": - api_base = ( - api_base - or get_secret("EMPOWER_API_BASE") - or "https://app.empower.dev/api/v1" - ) # type: ignore + api_base = api_base or get_secret("EMPOWER_API_BASE") or "https://app.empower.dev/api/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("EMPOWER_API_KEY") elif custom_llm_provider == "groq": ( api_base, dynamic_api_key, - ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "bedrock_mantle": ( api_base, dynamic_api_key, - ) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "nvidia_nim": # nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret("NVIDIA_NIM_API_BASE") - or "https://integrate.api.nvidia.com/v1" - ) # type: ignore + api_base = api_base or get_secret("NVIDIA_NIM_API_BASE") or "https://integrate.api.nvidia.com/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("NVIDIA_NIM_API_KEY") elif custom_llm_provider == "cerebras": - api_base = ( - api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("CEREBRAS_API_KEY") elif custom_llm_provider == "baseten": # Use BasetenConfig to determine the appropriate API base URL if api_base is None: api_base = litellm.BasetenConfig.get_api_base_for_model(model) else: - api_base = ( - api_base - or get_secret_str("BASETEN_API_BASE") - or "https://inference.baseten.co/v1" - ) + api_base = api_base or get_secret_str("BASETEN_API_BASE") or "https://inference.baseten.co/v1" dynamic_api_key = api_key or get_secret_str("BASETEN_API_KEY") elif custom_llm_provider == "sambanova": - api_base = ( - api_base - or get_secret("SAMBANOVA_API_BASE") - or "https://api.sambanova.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("SAMBANOVA_API_BASE") or "https://api.sambanova.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("SAMBANOVA_API_KEY") elif custom_llm_provider == "meta_llama": - api_base = ( - api_base - or get_secret("LLAMA_API_BASE") - or "https://api.llama.com/compat/v1" - ) # type: ignore + api_base = api_base or get_secret("LLAMA_API_BASE") or "https://api.llama.com/compat/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("LLAMA_API_KEY") elif custom_llm_provider == "nebius": - api_base = ( - api_base - or get_secret("NEBIUS_API_BASE") - or "https://api.studio.nebius.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("NEBIUS_API_BASE") or "https://api.studio.nebius.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") elif custom_llm_provider == "ollama": - api_base = ( - api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" - ) # type: ignore + api_base = api_base or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" # type: ignore dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") - elif (custom_llm_provider == "ai21_chat") or ( - custom_llm_provider == "ai21" and model in litellm.ai21_chat_models - ): - api_base = ( - api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" - ) # type: ignore + elif (custom_llm_provider == "ai21_chat") or (custom_llm_provider == "ai21" and model in litellm.ai21_chat_models): + api_base = api_base or get_secret("AI21_API_BASE") or "https://api.ai21.com/studio/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("AI21_API_KEY") custom_llm_provider = "ai21_chat" elif custom_llm_provider == "volcengine": # volcengine is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 - api_base = ( - api_base - or get_secret("VOLCENGINE_API_BASE") - or "https://ark.cn-beijing.volces.com/api/v3" - ) # type: ignore + api_base = api_base or get_secret("VOLCENGINE_API_BASE") or "https://ark.cn-beijing.volces.com/api/v3" # type: ignore dynamic_api_key = api_key or get_secret_str("VOLCENGINE_API_KEY") elif custom_llm_provider == "codestral": # codestral is openai compatible, we just need to set this to custom_openai and have the api_base be https://codestral.mistral.ai/v1 - api_base = ( - api_base - or get_secret("CODESTRAL_API_BASE") - or "https://codestral.mistral.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("CODESTRAL_API_BASE") or "https://codestral.mistral.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("CODESTRAL_API_KEY") elif custom_llm_provider == "hosted_vllm": # vllm is openai compatible, we just need to set this to custom_openai ( api_base, dynamic_api_key, - ) = litellm.HostedVLLMChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HostedVLLMChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "llamafile": # llamafile is OpenAI compatible. ( api_base, dynamic_api_key, - ) = litellm.LlamafileChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LlamafileChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "datarobot": # DataRobot is OpenAI compatible. ( api_base, dynamic_api_key, - ) = litellm.DataRobotConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DataRobotConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "lm_studio": # lm_studio is openai compatible, we just need to set this to custom_openai ( api_base, dynamic_api_key, - ) = litellm.LMStudioChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LMStudioChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "chuizi": # chuizi is openai compatible, unified AI gateway - api_base = ( - api_base - or get_secret("CHUIZI_API_BASE") - or "https://api.chuizi.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("CHUIZI_API_BASE") or "https://api.chuizi.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("CHUIZI_API_KEY") elif custom_llm_provider == "deepseek": # deepseek is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.deepseek.com/v1 - api_base = ( - api_base - or get_secret("DEEPSEEK_API_BASE") - or "https://api.deepseek.com/beta" - ) # type: ignore + api_base = api_base or get_secret("DEEPSEEK_API_BASE") or "https://api.deepseek.com/beta" # type: ignore dynamic_api_key = api_key or get_secret_str("DEEPSEEK_API_KEY") elif custom_llm_provider == "fireworks_ai": @@ -706,9 +589,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ( api_base, dynamic_api_key, - ) = litellm.FireworksAIConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.FireworksAIConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "azure_ai": ( api_base, @@ -728,45 +609,31 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ( api_base, dynamic_api_key, - ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "mistral": ( api_base, dynamic_api_key, - ) = litellm.MistralConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MistralConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "jina_ai": ( custom_llm_provider, api_base, dynamic_api_key, - ) = litellm.JinaAIEmbeddingConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.JinaAIEmbeddingConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "xai": ( api_base, dynamic_api_key, - ) = litellm.XAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.XAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "zai": ( api_base, dynamic_api_key, - ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "together_ai": - api_base = ( - api_base - or get_secret_str("TOGETHER_AI_API_BASE") - or "https://api.together.xyz/v1" - ) # type: ignore + api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1" # type: ignore dynamic_api_key = api_key or ( get_secret_str("TOGETHER_API_KEY") or get_secret_str("TOGETHER_AI_API_KEY") @@ -774,22 +641,10 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 or get_secret_str("TOGETHER_AI_TOKEN") ) elif custom_llm_provider == "friendliai": - api_base = ( - api_base - or get_secret("FRIENDLI_API_BASE") - or "https://api.friendli.ai/serverless/v1" - ) # type: ignore - dynamic_api_key = ( - api_key - or get_secret_str("FRIENDLIAI_API_KEY") - or get_secret_str("FRIENDLI_TOKEN") - ) + api_base = api_base or get_secret("FRIENDLI_API_BASE") or "https://api.friendli.ai/serverless/v1" # type: ignore + dynamic_api_key = api_key or get_secret_str("FRIENDLIAI_API_KEY") or get_secret_str("FRIENDLI_TOKEN") elif custom_llm_provider == "galadriel": - api_base = ( - api_base - or get_secret("GALADRIEL_API_BASE") - or "https://api.galadriel.com/v1" - ) # type: ignore + api_base = api_base or get_secret("GALADRIEL_API_BASE") or "https://api.galadriel.com/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("GALADRIEL_API_KEY") elif custom_llm_provider == "github_copilot": ( @@ -804,167 +659,115 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 api_base, dynamic_api_key, custom_llm_provider, - ) = litellm.ChatGPTConfig()._get_openai_compatible_provider_info( - model, api_base, api_key, custom_llm_provider - ) + ) = litellm.ChatGPTConfig()._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) elif custom_llm_provider == "novita": - api_base = ( - api_base - or get_secret("NOVITA_API_BASE") - or "https://api.novita.ai/v3/openai" - ) # type: ignore + api_base = api_base or get_secret("NOVITA_API_BASE") or "https://api.novita.ai/v3/openai" # type: ignore dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY") elif custom_llm_provider == "snowflake": ( api_base, dynamic_api_key, - ) = litellm.SnowflakeConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.SnowflakeConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "gradient_ai": ( api_base, dynamic_api_key, - ) = litellm.GradientAIConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.GradientAIConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "featherless_ai": ( api_base, dynamic_api_key, - ) = litellm.FeatherlessAIConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.FeatherlessAIConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "nscale": ( api_base, dynamic_api_key, - ) = litellm.NscaleConfig()._get_openai_compatible_provider_info( - api_base=api_base, api_key=api_key - ) + ) = litellm.NscaleConfig()._get_openai_compatible_provider_info(api_base=api_base, api_key=api_key) elif custom_llm_provider == "heroku": ( api_base, dynamic_api_key, - ) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "dashscope": ( api_base, dynamic_api_key, - ) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "moonshot": ( api_base, dynamic_api_key, - ) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MoonshotChatConfig()._get_openai_compatible_provider_info(api_base, api_key) # publicai is now handled by JSON config (see litellm/llms/openai_like/providers.json) elif custom_llm_provider == "docker_model_runner": ( api_base, dynamic_api_key, - ) = litellm.DockerModelRunnerChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.DockerModelRunnerChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "v0": ( api_base, dynamic_api_key, - ) = litellm.V0ChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.V0ChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "morph": ( api_base, dynamic_api_key, - ) = litellm.MorphChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.MorphChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "lambda_ai": ( api_base, dynamic_api_key, - ) = litellm.LambdaAIChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LambdaAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "hyperbolic": ( api_base, dynamic_api_key, - ) = litellm.HyperbolicChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.HyperbolicChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "vercel_ai_gateway": ( api_base, dynamic_api_key, - ) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.VercelAIGatewayConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "aiml": ( api_base, dynamic_api_key, - ) = litellm.AIMLChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.AIMLChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "wandb": - api_base = ( - api_base - or get_secret("WANDB_API_BASE") - or "https://api.inference.wandb.ai/v1" - ) # type: ignore + api_base = api_base or get_secret("WANDB_API_BASE") or "https://api.inference.wandb.ai/v1" # type: ignore dynamic_api_key = api_key or get_secret_str("WANDB_API_KEY") elif custom_llm_provider == "lemonade": ( api_base, dynamic_api_key, - ) = litellm.LemonadeChatConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.LemonadeChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "clarifai": ( api_base, dynamic_api_key, - ) = litellm.ClarifaiConfig()._get_openai_compatible_provider_info( - api_base, api_key - ) + ) = litellm.ClarifaiConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "ragflow": full_model = f"ragflow/{model}" ( api_base, dynamic_api_key, _, - ) = litellm.RAGFlowConfig()._get_openai_compatible_provider_info( - full_model, api_base, api_key, "ragflow" - ) + ) = litellm.RAGFlowConfig()._get_openai_compatible_provider_info(full_model, api_base, api_key, "ragflow") model = full_model elif custom_llm_provider == "langgraph": # LangGraph is a custom provider, just need to set api_base - api_base = ( - api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" - ) + api_base = api_base or get_secret_str("LANGGRAPH_API_BASE") or "http://localhost:2024" dynamic_api_key = api_key or get_secret_str("LANGGRAPH_API_KEY") elif custom_llm_provider == "manus": # Manus is OpenAI compatible for responses API - api_base = ( - api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" - ) + api_base = api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im" dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY") if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): - raise Exception( - "dynamic_api_key needs to be a string. dynamic_api_key={}".format( - dynamic_api_key - ) - ) + raise Exception("dynamic_api_key needs to be a string. dynamic_api_key={}".format(dynamic_api_key)) if dynamic_api_key is None and api_key is not None: dynamic_api_key = api_key return model, custom_llm_provider, dynamic_api_key, api_base diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 1535ce513cb..41caddfdb91 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -9,13 +9,12 @@ import io from unittest.mock import patch -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import pytest import litellm from litellm.types.router import LiteLLM_Params + def test_get_llm_provider(): _, response, _, _ = litellm.get_llm_provider(model="anthropic.claude-v2:1") @@ -26,9 +25,7 @@ def test_get_llm_provider(): def test_get_llm_provider_fireworks(): # tests finetuned fireworks models - https://github.com/BerriAI/litellm/issues/4923 - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model="fireworks_ai/accounts/my-test-1234" - ) + model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="fireworks_ai/accounts/my-test-1234") assert custom_llm_provider == "fireworks_ai" assert model == "accounts/my-test-1234" @@ -52,10 +49,7 @@ def test_get_llm_provider_mistral_custom_api_base(): ) assert custom_llm_provider == "mistral" assert model == "mistral-large-fr" - assert ( - api_base - == "https://mistral-large-fr-ishaan.francecentral.inference.ai.azure.com/v1" - ) + assert api_base == "https://mistral-large-fr-ishaan.francecentral.inference.ai.azure.com/v1" def test_get_llm_provider_deepseek_custom_api_base(): @@ -118,7 +112,6 @@ def test_get_llm_provider_cohere_chat_test2(): def test_get_llm_provider_azure_o1(): - model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( model="azure/o1-mini", ) @@ -150,9 +143,7 @@ def test_default_api_base(): continue for other_provider in litellm.provider_list: - if other_provider != provider and provider != "{}_chat".format( - other_provider.value - ): + if other_provider != provider and provider != "{}_chat".format(other_provider.value): if provider == "codestral" and other_provider == "mistral": continue elif provider == "github" and other_provider == "azure": @@ -244,16 +235,16 @@ def test_xai_api_base(model): "api_key": "xai-my-specialkey", "litellm_params": None, } - model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( - **args - ) + model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider(**args) assert custom_llm_provider == "xai" assert model == "grok-2-vision-latest" assert api_base == "https://api.x.ai/v1" assert dynamic_api_key == "xai-my-specialkey" + # -------- Tests for force_use_litellm_proxy --------- + def test_get_litellm_proxy_custom_llm_provider(): """ Tests force_use_litellm_proxy uses LITELLM_PROXY_API_BASE and LITELLM_PROXY_API_KEY from env. @@ -262,17 +253,19 @@ def test_get_litellm_proxy_custom_llm_provider(): expected_api_base = "http://localhost:8000" expected_api_key = "test_proxy_key" - with patch.dict(os.environ, { - "LITELLM_PROXY_API_BASE": expected_api_base, - "LITELLM_PROXY_API_KEY": expected_api_key - }, clear=True): - model, provider, key, base = litellm.LiteLLMProxyChatConfig().litellm_proxy_get_custom_llm_provider_info(model=test_model) + with patch.dict( + os.environ, {"LITELLM_PROXY_API_BASE": expected_api_base, "LITELLM_PROXY_API_KEY": expected_api_key}, clear=True + ): + model, provider, key, base = litellm.LiteLLMProxyChatConfig().litellm_proxy_get_custom_llm_provider_info( + model=test_model + ) assert model == test_model assert provider == "litellm_proxy" assert key == expected_api_key assert base == expected_api_base + def test_get_litellm_proxy_with_args_override_env_vars(): """ Tests force_use_litellm_proxy uses api_base and api_key args over environment variables. @@ -280,18 +273,15 @@ def test_get_litellm_proxy_with_args_override_env_vars(): test_model = "gpt-4" arg_api_base = "http://custom-proxy.com" arg_api_key = "custom_key_from_arg" - + env_api_base = "http://env-proxy.com" env_api_key = "env_key" - with patch.dict(os.environ, { - "LITELLM_PROXY_API_BASE": env_api_base, - "LITELLM_PROXY_API_KEY": env_api_key - }, clear=True): + with patch.dict( + os.environ, {"LITELLM_PROXY_API_BASE": env_api_base, "LITELLM_PROXY_API_KEY": env_api_key}, clear=True + ): model, provider, key, base = litellm.LiteLLMProxyChatConfig().litellm_proxy_get_custom_llm_provider_info( - model=test_model, - api_base=arg_api_base, - api_key=arg_api_key + model=test_model, api_base=arg_api_base, api_key=arg_api_key ) assert model == test_model @@ -299,6 +289,7 @@ def test_get_litellm_proxy_with_args_override_env_vars(): assert key == arg_api_key assert base == arg_api_base + def test_get_litellm_proxy_model_prefix_stripping(): """ Tests force_use_litellm_proxy strips 'litellm_proxy/' prefix from model name. @@ -308,19 +299,22 @@ def test_get_litellm_proxy_model_prefix_stripping(): expected_api_base = "http://localhost:4000" expected_api_key = "proxy_secret_key" - with patch.dict(os.environ, { - "LITELLM_PROXY_API_BASE": expected_api_base, - "LITELLM_PROXY_API_KEY": expected_api_key - }, clear=True): - model, provider, key, base = litellm.LiteLLMProxyChatConfig().litellm_proxy_get_custom_llm_provider_info(model=original_model) + with patch.dict( + os.environ, {"LITELLM_PROXY_API_BASE": expected_api_base, "LITELLM_PROXY_API_KEY": expected_api_key}, clear=True + ): + model, provider, key, base = litellm.LiteLLMProxyChatConfig().litellm_proxy_get_custom_llm_provider_info( + model=original_model + ) assert model == expected_model assert provider == "litellm_proxy" assert key == expected_api_key assert base == expected_api_base + # -------- Tests for get_llm_provider triggering use_litellm_proxy --------- + def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true(): """ Tests get_llm_provider uses litellm_proxy when USE_LITELLM_PROXY is "True". @@ -330,13 +324,13 @@ def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true(): proxy_api_base = "http://my-global-proxy.com" proxy_api_key = "global_proxy_key" - with patch.dict(os.environ, { - "USE_LITELLM_PROXY": "True", - "LITELLM_PROXY_API_BASE": proxy_api_base, - "LITELLM_PROXY_API_KEY": proxy_api_key - }, clear=True): + with patch.dict( + os.environ, + {"USE_LITELLM_PROXY": "True", "LITELLM_PROXY_API_BASE": proxy_api_base, "LITELLM_PROXY_API_KEY": proxy_api_key}, + clear=True, + ): model, provider, key, base = litellm.get_llm_provider(model=test_model_input) - + print("get_llm_provider", model, provider, key, base) assert model == expected_model_output @@ -344,6 +338,7 @@ def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true(): assert key == proxy_api_key assert base == proxy_api_base + def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true_model_prefix(): """ Tests get_llm_provider with USE_LITELLM_PROXY="True" and model prefix "litellm_proxy/". @@ -353,11 +348,11 @@ def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true_model_prefix(): proxy_api_base = "http://another-proxy.net" proxy_api_key = "another_key" - with patch.dict(os.environ, { - "USE_LITELLM_PROXY": "True", - "LITELLM_PROXY_API_BASE": proxy_api_base, - "LITELLM_PROXY_API_KEY": proxy_api_key - }, clear=True): + with patch.dict( + os.environ, + {"USE_LITELLM_PROXY": "True", "LITELLM_PROXY_API_BASE": proxy_api_base, "LITELLM_PROXY_API_KEY": proxy_api_key}, + clear=True, + ): model, provider, key, base = litellm.get_llm_provider(model=test_model_input) assert model == expected_model_output @@ -371,18 +366,16 @@ def test_get_llm_provider_use_proxy_arg_true(): Tests get_llm_provider uses litellm_proxy when use_proxy=True argument is passed. """ test_model_input = "mistral/mistral-large" - expected_model_output = "mistral/mistral-large" # force_use_litellm_proxy keep the model name + expected_model_output = "mistral/mistral-large" # force_use_litellm_proxy keep the model name proxy_api_base = "http://my-arg-proxy.com" proxy_api_key = "arg_proxy_key" - + # Ensure LITELLM_PROXY_ALWAYS is not set or False - with patch.dict(os.environ, { - "LITELLM_PROXY_API_BASE": proxy_api_base, - "LITELLM_PROXY_API_KEY": proxy_api_key - }, clear=True): # clear=True removes LITELLM_PROXY_ALWAYS if it was set by other tests + with patch.dict( + os.environ, {"LITELLM_PROXY_API_BASE": proxy_api_base, "LITELLM_PROXY_API_KEY": proxy_api_key}, clear=True + ): # clear=True removes LITELLM_PROXY_ALWAYS if it was set by other tests model, provider, key, base = litellm.get_llm_provider( - model=test_model_input, - litellm_params=LiteLLM_Params(use_litellm_proxy=True, model=test_model_input) + model=test_model_input, litellm_params=LiteLLM_Params(use_litellm_proxy=True, model=test_model_input) ) assert model == expected_model_output @@ -390,6 +383,7 @@ def test_get_llm_provider_use_proxy_arg_true(): assert key == proxy_api_key assert base == proxy_api_base + def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): """ Tests get_llm_provider with use_proxy=True and explicit api_base/api_key args. @@ -397,7 +391,7 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): """ test_model_input = "anthropic/claude-3-opus" expected_model_output = "anthropic/claude-3-opus" - + arg_api_base = "http://specific-proxy-endpoint.org" arg_api_key = "specific_key_for_call" @@ -405,21 +399,22 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): env_proxy_api_base = "http://env-default-proxy.com" env_proxy_api_key = "env_default_key" - with patch.dict(os.environ, { - "LITELLM_PROXY_API_BASE": env_proxy_api_base, - "LITELLM_PROXY_API_KEY": env_proxy_api_key - }, clear=True): + with patch.dict( + os.environ, + {"LITELLM_PROXY_API_BASE": env_proxy_api_base, "LITELLM_PROXY_API_KEY": env_proxy_api_key}, + clear=True, + ): model, provider, key, base = litellm.get_llm_provider( - model=test_model_input, + model=test_model_input, api_base=arg_api_base, api_key=arg_api_key, - litellm_params=LiteLLM_Params(use_litellm_proxy=True, model=test_model_input) + litellm_params=LiteLLM_Params(use_litellm_proxy=True, model=test_model_input), ) assert model == expected_model_output assert provider == "litellm_proxy" assert key == arg_api_key # Should use the argument key - assert base == arg_api_base # Should use the argument base + assert base == arg_api_base # Should use the argument base def test_get_llm_provider_chuizi():