diff --git a/litellm/main.py b/litellm/main.py index 80a9a94a34..cf37d46d15 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -943,6 +943,7 @@ def completion( output_cost_per_token=output_cost_per_token, cooldown_time=cooldown_time, text_completion=kwargs.get("text_completion"), + azure_ad_token_provider=kwargs.get("azure_ad_token_provider"), ) logging.update_environment_variables( model=model, @@ -3229,6 +3230,7 @@ def embedding( "model_config", "cooldown_time", "tags", + "azure_ad_token_provider", ] default_params = openai_params + litellm_params non_default_params = { diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 65c7f70525..7d9b97246c 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,9 +1,11 @@ model_list: - model_name: gpt-4 litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ + model: azure/chatgpt-v-2 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + azure_ad_token_provider: True + guardrails: - guardrail_name: "lakera-pre-guard" diff --git a/litellm/router_utils/client_initalization_utils.py b/litellm/router_utils/client_initalization_utils.py index f396defb51..b0b13fac7f 100644 --- a/litellm/router_utils/client_initalization_utils.py +++ b/litellm/router_utils/client_initalization_utils.py @@ -1,7 +1,7 @@ import asyncio import os import traceback -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable import httpx import openai @@ -173,6 +173,11 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): organization = litellm.get_secret(organization_env_name) litellm_params["organization"] = organization + azure_ad_token_provider = litellm_params.get("azure_ad_token_provider", None) + if azure_ad_token_provider is not None: + verbose_router_logger.debug("Using Azure AD Token Provider for Azure Auth") + azure_ad_token_provider = get_azure_ad_token_from_entrata_id() + if custom_llm_provider == "azure" or custom_llm_provider == "azure_text": if api_base is None or not isinstance(api_base, str): filtered_litellm_params = { @@ -190,7 +195,9 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): if azure_ad_token.startswith("oidc/"): azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token) if api_version is None: - api_version = os.getenv("AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION) + api_version = os.getenv( + "AZURE_API_VERSION", litellm.AZURE_DEFAULT_API_VERSION + ) if "gateway.ai.cloudflare.com" in api_base: if not api_base.endswith("/"): @@ -304,6 +311,11 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): "api_version": api_version, "azure_ad_token": azure_ad_token, } + + if azure_ad_token_provider is not None: + azure_client_params["azure_ad_token_provider"] = ( + azure_ad_token_provider + ) from litellm.llms.azure import select_azure_base_url_or_endpoint # this decides if we should set azure_endpoint or base_url on Azure OpenAI Client @@ -493,3 +505,17 @@ def set_client(litellm_router_instance: LitellmRouter, model: dict): ttl=client_ttl, local_only=True, ) # cache for 1 hr + + +def get_azure_ad_token_from_entrata_id() -> Callable[[], str]: + from azure.identity import DefaultAzureCredential, get_bearer_token_provider + + verbose_router_logger.debug("Getting Azure AD Token from Entrata ID") + + token_provider = get_bearer_token_provider( + DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default" + ) + + verbose_router_logger.debug("token_provider %s", token_provider) + + return token_provider