diff --git a/docs/my-website/docs/providers/litellm_proxy.md b/docs/my-website/docs/providers/litellm_proxy.md index a66423dac5..bebaf1230f 100644 --- a/docs/my-website/docs/providers/litellm_proxy.md +++ b/docs/my-website/docs/providers/litellm_proxy.md @@ -155,6 +155,51 @@ response = litellm.rerank( api_key="your-litellm-proxy-api-key" ) ``` -## **Usage with Langchain, LLamaindex, OpenAI Js, Anthropic SDK, Instructor** -#### [Follow this doc to see how to use litellm proxy with langchain, llamaindex, anthropic etc](../proxy/user_keys) \ No newline at end of file + +## Integration with Other Libraries + +LiteLLM Proxy works seamlessly with Langchain, LlamaIndex, OpenAI JS, Anthropic SDK, Instructor, and more. + +[Learn how to use LiteLLM proxy with these libraries →](../proxy/user_keys) + +## Flags to send requests to litellm proxy + +Use the following options to route all requests through your LiteLLM proxy, regardless of the model specified. + +When enabled, requests will use `LITELLM_PROXY_API_BASE` with `LITELLM_PROXY_API_KEY` as the authentication. + +### Option 1: Set Globally in Code + +```python +# Set the flag globally for all requests +litellm.use_litellm_proxy = True + +response = litellm.completion( + model="vertex_ai/gemini-2.0-flash-001", + messages=[{"role": "user", "content": "Hello, how are you?"}] +) +``` + +### Option 2: Control via Environment Variable + +```python +# Control proxy usage through environment variable +os.environ["USE_LITELLM_PROXY"] = "True" + +response = litellm.completion( + model="vertex_ai/gemini-2.0-flash-001", + messages=[{"role": "user", "content": "Hello, how are you?"}] +) +``` + +### Option 3: Set Per Request + +```python +# Enable proxy for specific requests only +response = litellm.completion( + model="vertex_ai/gemini-2.0-flash-001", + messages=[{"role": "user", "content": "Hello, how are you?"}], + use_litellm_proxy=True +) +``` diff --git a/litellm/__init__.py b/litellm/__init__.py index d8ec5a747c..f1b481612c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -132,22 +132,22 @@ prometheus_initialize_budget_metrics: Optional[bool] = False require_auth_for_metrics_endpoint: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload -gcs_pub_sub_use_v1: Optional[ - bool -] = False # if you want to use v1 gcs pubsub logged payload -generic_api_use_v1: Optional[ - bool -] = False # if you want to use v1 generic api logged payload +gcs_pub_sub_use_v1: Optional[bool] = ( + False # if you want to use v1 gcs pubsub logged payload +) +generic_api_use_v1: Optional[bool] = ( + False # if you want to use v1 generic api logged payload +) argilla_transformation_object: Optional[Dict[str, Any]] = None -_async_input_callback: List[ - Union[str, Callable, CustomLogger] -] = [] # internal variable - async custom callbacks are routed here. -_async_success_callback: List[ - Union[str, Callable, CustomLogger] -] = [] # internal variable - async custom callbacks are routed here. -_async_failure_callback: List[ - Union[str, Callable, CustomLogger] -] = [] # internal variable - async custom callbacks are routed here. +_async_input_callback: List[Union[str, Callable, CustomLogger]] = ( + [] +) # internal variable - async custom callbacks are routed here. +_async_success_callback: List[Union[str, Callable, CustomLogger]] = ( + [] +) # internal variable - async custom callbacks are routed here. +_async_failure_callback: List[Union[str, Callable, CustomLogger]] = ( + [] +) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] post_call_rules: List[Callable] = [] turn_off_message_logging: Optional[bool] = False @@ -155,18 +155,18 @@ log_raw_request_response: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False filter_invalid_headers: Optional[bool] = False -add_user_information_to_llm_headers: Optional[ - bool -] = None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers +add_user_information_to_llm_headers: Optional[bool] = ( + None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers +) store_audit_logs = False # Enterprise feature, allow users to see audit logs ### end of callbacks ############# -email: Optional[ - str -] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -token: Optional[ - str -] = None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +email: Optional[str] = ( + None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +) +token: Optional[str] = ( + None # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +) telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) @@ -204,6 +204,9 @@ common_cloud_provider_auth_params: dict = { "params": ["project", "region_name", "token"], "providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"], } +use_litellm_proxy: bool = ( + False # when True, requests will be sent to the specified litellm proxy endpoint +) use_client: bool = False ssl_verify: Union[str, bool] = True ssl_certificate: Optional[str] = None @@ -243,20 +246,24 @@ enable_loadbalancing_on_batch_endpoints: Optional[bool] = None enable_caching_on_provider_specific_optional_params: bool = ( False # feature-flag for caching on optional params - e.g. 'top_k' ) -caching: bool = False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -caching_with_models: bool = False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 -cache: Optional[ - Cache -] = None # cache object <- use this - https://docs.litellm.ai/docs/caching +caching: bool = ( + False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +) +caching_with_models: bool = ( + False # # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 +) +cache: Optional[Cache] = ( + None # cache object <- use this - https://docs.litellm.ai/docs/caching +) default_in_memory_ttl: Optional[float] = None default_redis_ttl: Optional[float] = None default_redis_batch_cache_expiry: Optional[float] = None model_alias_map: Dict[str, str] = {} model_group_alias_map: Dict[str, str] = {} max_budget: float = 0.0 # set the max budget across all providers -budget_duration: Optional[ - str -] = None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). +budget_duration: Optional[str] = ( + None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). +) default_soft_budget: float = ( DEFAULT_SOFT_BUDGET # by default all litellm proxy keys have a soft budget of 50.0 ) @@ -265,11 +272,15 @@ forward_traceparent_to_llm_provider: bool = False _current_cost = 0.0 # private variable, used if max budget is set error_logs: Dict = {} -add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt +add_function_to_prompt: bool = ( + False # if function calling not supported by api, append function call details to system prompt +) client_session: Optional[httpx.Client] = None aclient_session: Optional[httpx.AsyncClient] = None model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks' -model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +model_cost_map_url: str = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" +) suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None @@ -292,7 +303,9 @@ disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None custom_prometheus_metadata_labels: List[str] = [] #### REQUEST PRIORITIZATION #### priority_reservation: Optional[Dict[str, float]] = None -force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. +force_ipv4: bool = ( + False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6. +) module_level_aclient = AsyncHTTPHandler( timeout=request_timeout, client_alias="module level aclient" ) @@ -306,13 +319,13 @@ fallbacks: Optional[List] = None context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None allowed_fails: int = 3 -num_retries_per_request: Optional[ - int -] = None # for the request overall (incl. fallbacks + model retries) +num_retries_per_request: Optional[int] = ( + None # for the request overall (incl. fallbacks + model retries) +) ####### SECRET MANAGERS ##################### -secret_manager_client: Optional[ - Any -] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. +secret_manager_client: Optional[Any] = ( + None # list of instantiated key management clients - e.g. azure kv, infisical, etc. +) _google_kms_resource_name: Optional[str] = None _key_management_system: Optional[KeyManagementSystem] = None _key_management_settings: KeyManagementSettings = KeyManagementSettings() @@ -1092,10 +1105,10 @@ from .types.llms.custom_llm import CustomLLMItem from .types.utils import GenericStreamingChunk custom_provider_map: List[CustomLLMItem] = [] -_custom_providers: List[ - str -] = [] # internal helper util, used to track names of custom providers -disable_hf_tokenizer_download: Optional[ - bool -] = None # disable huggingface tokenizer download. Defaults to openai clk100 +_custom_providers: List[str] = ( + [] +) # internal helper util, used to track names of custom providers +disable_hf_tokenizer_download: Optional[bool] = ( + None # disable huggingface tokenizer download. Defaults to openai clk100 +) global_disable_no_log_param: bool = False diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 7e8f60bd4f..19c8ec8d80 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -59,6 +59,7 @@ def get_litellm_params( async_call: Optional[bool] = None, ssl_verify: Optional[bool] = None, merge_reasoning_content_in_choices: Optional[bool] = None, + use_litellm_proxy: Optional[bool] = None, api_version: Optional[str] = None, max_retries: Optional[int] = None, **kwargs, @@ -115,5 +116,6 @@ def get_litellm_params( "bucket_name": kwargs.get("bucket_name"), "vertex_credentials": kwargs.get("vertex_credentials"), "vertex_project": kwargs.get("vertex_project"), + "use_litellm_proxy": use_litellm_proxy, } return litellm_params diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 29f6b2d3f6..6d03964f2b 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -102,8 +102,15 @@ def get_llm_provider( # noqa: PLR0915 Return model, custom_llm_provider, dynamic_api_key, api_base """ try: + 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 + ) + ## IF LITELLM PARAMS GIVEN ## - if litellm_params is not None: + if litellm_params: assert ( custom_llm_provider is None and api_base is None and api_key is None ), "Either pass in litellm_params or the custom_llm_provider/api_base/api_key. Otherwise, these values will be overriden." @@ -538,8 +545,12 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) dynamic_api_key = api_key or get_secret_str("GITHUB_API_KEY") elif custom_llm_provider == "litellm_proxy": - api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") - dynamic_api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY") + ( + api_base, + dynamic_api_key, + ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info( + api_base=api_base, api_key=api_key + ) elif custom_llm_provider == "mistral": ( diff --git a/litellm/llms/litellm_proxy/chat/transformation.py b/litellm/llms/litellm_proxy/chat/transformation.py index 22013198ba..6896b37e61 100644 --- a/litellm/llms/litellm_proxy/chat/transformation.py +++ b/litellm/llms/litellm_proxy/chat/transformation.py @@ -4,17 +4,18 @@ Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions` from typing import List, Optional, Tuple -from litellm.secret_managers.main import get_secret_str +from litellm.secret_managers.main import get_secret_bool, get_secret_str +from litellm.types.router import LiteLLM_Params from ...openai.chat.gpt_transformation import OpenAIGPTConfig class LiteLLMProxyChatConfig(OpenAIGPTConfig): def get_supported_openai_params(self, model: str) -> List: - list = super().get_supported_openai_params(model) - list.append("thinking") - list.append("reasoning_effort") - return list + params_list = super().get_supported_openai_params(model) + params_list.append("thinking") + params_list.append("reasoning_effort") + return params_list def _map_openai_params( self, @@ -52,3 +53,63 @@ class LiteLLMProxyChatConfig(OpenAIGPTConfig): @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: return api_key or get_secret_str("LITELLM_PROXY_API_KEY") + + @staticmethod + def _should_use_litellm_proxy_by_default( + litellm_params: Optional[LiteLLM_Params] = None, + ): + """ + Returns True if litellm proxy should be used by default for a given request + + Issue: https://github.com/BerriAI/litellm/issues/10559 + + Use case: + - When using Google ADK, users want a flag to dynamically enable sending the request to litellm proxy or not + - Allow the model name to be passed in original format and still use litellm proxy: + "gemini/gemini-1.5-pro", "openai/gpt-4", "mistral/llama-2-70b-chat" etc. + """ + import litellm + + if get_secret_bool("USE_LITELLM_PROXY") is True: + return True + if litellm_params and litellm_params.use_litellm_proxy is True: + return True + if litellm.use_litellm_proxy is True: + return True + return False + + @staticmethod + def litellm_proxy_get_custom_llm_provider_info( + model: str, api_base: Optional[str] = None, api_key: Optional[str] = None + ) -> Tuple[str, str, Optional[str], Optional[str]]: + """ + Force use litellm proxy for all models + + Issue: https://github.com/BerriAI/litellm/issues/10559 + + Expected behavior: + - custom_llm_provider will be 'litellm_proxy' + - api_base = api_base OR LITELLM_PROXY_API_BASE + - api_key = api_key OR LITELLM_PROXY_API_KEY + + Use case: + - When using Google ADK, users want a flag to dynamically enable sending the request to litellm proxy or not + - Allow the model name to be passed in original format and still use litellm proxy: + "gemini/gemini-1.5-pro", "openai/gpt-4", "mistral/llama-2-70b-chat" etc. + + Return model, custom_llm_provider, dynamic_api_key, api_base + """ + import litellm + + custom_llm_provider = "litellm_proxy" + if model.startswith("litellm_proxy/"): + model = model.split("/", 1)[1] + + ( + api_base, + api_key, + ) = litellm.LiteLLMProxyChatConfig()._get_openai_compatible_provider_info( + api_base=api_base, api_key=api_key + ) + + return model, custom_llm_provider, api_key, api_base diff --git a/litellm/main.py b/litellm/main.py index f8c3ec9b98..e09c8b3418 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1221,6 +1221,7 @@ def completion( # type: ignore # noqa: PLR0915 merge_reasoning_content_in_choices=kwargs.get( "merge_reasoning_content_in_choices", None ), + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), api_version=api_version, azure_ad_token=kwargs.get("azure_ad_token"), tenant_id=kwargs.get("tenant_id"), @@ -2725,9 +2726,9 @@ def completion( # type: ignore # noqa: PLR0915 "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None ): - optional_params[ - "aws_region_name" - ] = aws_bedrock_client.meta.region_name + optional_params["aws_region_name"] = ( + aws_bedrock_client.meta.region_name + ) bedrock_route = BedrockModelInfo.get_bedrock_route(model) if bedrock_route == "converse": @@ -4451,9 +4452,9 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[ - Union[BaseModel, AdapterCompletionStreamWrapper] - ] = None + translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( + None + ) if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params( response=response @@ -5920,9 +5921,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(content_chunks) > 0: - response["choices"][0]["message"][ - "content" - ] = processor.get_combined_content(content_chunks) + response["choices"][0]["message"]["content"] = ( + processor.get_combined_content(content_chunks) + ) reasoning_chunks = [ chunk @@ -5933,9 +5934,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"][ - "reasoning_content" - ] = processor.get_combined_reasoning_content(reasoning_chunks) + response["choices"][0]["message"]["reasoning_content"] = ( + processor.get_combined_reasoning_content(reasoning_chunks) + ) audio_chunks = [ chunk diff --git a/litellm/types/router.py b/litellm/types/router.py index bd2990d5fb..f563467135 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -96,16 +96,18 @@ class ModelInfo(BaseModel): id: Optional[ str ] # Allow id to be optional on input, but it will always be present as a str in the model instance - db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. + db_model: bool = ( + False # used for proxy - to separate models which are stored in the db vs. config. + ) updated_at: Optional[datetime.datetime] = None updated_by: Optional[str] = None created_at: Optional[datetime.datetime] = None created_by: Optional[str] = None - base_model: Optional[ - str - ] = None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking + base_model: Optional[str] = ( + None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking + ) tier: Optional[Literal["free", "paid"]] = None """ @@ -180,12 +182,12 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): custom_llm_provider: Optional[str] = None tpm: Optional[int] = None rpm: Optional[int] = None - timeout: Optional[ - Union[float, str, httpx.Timeout] - ] = None # if str, pass in as os.environ/ - stream_timeout: Optional[ - Union[float, str] - ] = None # timeout when making stream=True calls, if str, pass in as os.environ/ + timeout: Optional[Union[float, str, httpx.Timeout]] = ( + None # if str, pass in as os.environ/ + ) + stream_timeout: Optional[Union[float, str]] = ( + None # timeout when making stream=True calls, if str, pass in as os.environ/ + ) max_retries: Optional[int] = None organization: Optional[str] = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None @@ -200,6 +202,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): max_budget: Optional[float] = None budget_duration: Optional[str] = None use_in_pass_through: Optional[bool] = False + use_litellm_proxy: Optional[bool] = False model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: Optional[bool] = False model_info: Optional[Dict] = None @@ -242,6 +245,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): budget_duration: Optional[str] = None, # Pass through params use_in_pass_through: Optional[bool] = False, + # Dynamic param to force using litellm proxy + use_litellm_proxy: Optional[bool] = False, # This will merge the reasoning content in the choices merge_reasoning_content_in_choices: Optional[bool] = False, model_info: Optional[Dict] = None, @@ -255,9 +260,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): if max_retries is not None and isinstance(max_retries, str): max_retries = int(max_retries) # cast to int # We need to keep max_retries in args since it's a parameter of GenericLiteLLMParams - args[ - "max_retries" - ] = max_retries # Put max_retries back in args after popping it + args["max_retries"] = ( + max_retries # Put max_retries back in args after popping it + ) super().__init__(**args, **params) def __contains__(self, key): @@ -312,6 +317,7 @@ class LiteLLM_Params(GenericLiteLLMParams): max_file_size_mb: Optional[float] = None, # will use deployment on pass-through endpoints if True use_in_pass_through: Optional[bool] = False, + use_litellm_proxy: Optional[bool] = False, **params, ): args = locals() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d4e4ab4041..bcc4bf7ab8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2041,6 +2041,7 @@ all_litellm_params = [ "litellm_credential_name", "allowed_openai_params", "litellm_session_id", + "use_litellm_proxy", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index b5ea128971..d536310cee 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -13,9 +13,8 @@ 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") @@ -244,3 +243,172 @@ def test_xai_api_base(model): 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. + """ + test_model = "gpt-3.5-turbo" + 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) + + 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. + """ + 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): + 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 + ) + + assert model == test_model + assert provider == "litellm_proxy" + 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. + """ + original_model = "litellm_proxy/claude-2" + expected_model = "claude-2" + 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) + + 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". + """ + test_model_input = "openai/gpt-4" + expected_model_output = "openai/gpt-4" + 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): + 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 + assert provider == "litellm_proxy" + 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/". + """ + test_model_input = "litellm_proxy/gpt-4-turbo" + expected_model_output = "gpt-4-turbo" + 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): + model, provider, key, base = litellm.get_llm_provider(model=test_model_input) + + assert model == expected_model_output + assert provider == "litellm_proxy" + assert key == proxy_api_key + assert base == proxy_api_base + + +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 + 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 + model, provider, key, base = litellm.get_llm_provider( + 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 == 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. + These args should be passed to force_use_litellm_proxy and override env vars. + """ + 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" + + # Set some env vars to ensure they are overridden + 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): + model, provider, key, base = litellm.get_llm_provider( + 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) + ) + + 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