phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(azure.py): fix azure client init

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-11 18:05:11 -07:00
parent c5a80cc690
commit 23ca233f42
2 changed files with 11 additions and 97 deletions
Download patch file Download diff file Expand all files Collapse all files

64
litellm/llms/azure/azure.py
View file

@ -261,6 +261,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
max_retries = DEFAULT_MAX_RETRIES max_retries = DEFAULT_MAX_RETRIES
json_mode: Optional[bool] = optional_params.pop("json_mode", False) json_mode: Optional[bool] = optional_params.pop("json_mode", False)
azure_client_params = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
api_key=api_key,
api_base=api_base,
model_name=model,
api_version=api_version,
)
### CHECK IF CLOUDFLARE AI GATEWAY ### ### CHECK IF CLOUDFLARE AI GATEWAY ###
### if so - set the model as part of the base url ### if so - set the model as part of the base url
if "gateway.ai.cloudflare.com" in api_base: if "gateway.ai.cloudflare.com" in api_base:
@ -321,6 +328,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
timeout=timeout, timeout=timeout,
client=client, client=client,
max_retries=max_retries, max_retries=max_retries,
azure_client_params=azure_client_params,
) )
else: else:
return self.acompletion( return self.acompletion(
@ -338,7 +346,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
logging_obj=logging_obj, logging_obj=logging_obj,
max_retries=max_retries, max_retries=max_retries,
convert_tool_call_to_json_mode=json_mode, convert_tool_call_to_json_mode=json_mode,
litellm_params=litellm_params, azure_client_params=azure_client_params,
) )
elif "stream" in optional_params and optional_params["stream"] is True: elif "stream" in optional_params and optional_params["stream"] is True:
return self.streaming( return self.streaming(
@ -375,28 +383,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
status_code=422, message="max retries must be an int" status_code=422, message="max retries must be an int"
) )
# init AzureOpenAI Client # init AzureOpenAI Client
azure_client_params = {
"api_version": api_version,
"azure_endpoint": api_base,
"azure_deployment": model,
"http_client": litellm.client_session,
"max_retries": max_retries,
"timeout": timeout,
}
azure_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
if azure_ad_token.startswith("oidc/"):
azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token)
azure_client_params["azure_ad_token"] = azure_ad_token
elif azure_ad_token_provider is not None:
azure_client_params["azure_ad_token_provider"] = (
azure_ad_token_provider
)
if ( if (
client is None client is None
or not isinstance(client, AzureOpenAI) or not isinstance(client, AzureOpenAI)
@ -467,19 +453,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
azure_ad_token_provider: Optional[Callable] = None, azure_ad_token_provider: Optional[Callable] = None,
convert_tool_call_to_json_mode: Optional[bool] = None, convert_tool_call_to_json_mode: Optional[bool] = None,
client=None, # this is the AsyncAzureOpenAI client=None, # this is the AsyncAzureOpenAI
litellm_params: Optional[dict] = None, azure_client_params: dict = {},
): ):
response = None response = None
try: try:
# init AzureOpenAI Client
azure_client_params = self.initialize_azure_sdk_client(
litellm_params=litellm_params or {},
api_key=api_key,
api_base=api_base,
model_name=model,
api_version=api_version,
)
# setting Azure client # setting Azure client
if client is None or dynamic_params: if client is None or dynamic_params:
azure_client = AsyncAzureOpenAI(**azure_client_params) azure_client = AsyncAzureOpenAI(**azure_client_params)
@ -636,28 +613,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
azure_ad_token: Optional[str] = None, azure_ad_token: Optional[str] = None,
azure_ad_token_provider: Optional[Callable] = None, azure_ad_token_provider: Optional[Callable] = None,
client=None, client=None,
azure_client_params: dict = {},
): ):
try: try:
# init AzureOpenAI Client
azure_client_params = {
"api_version": api_version,
"azure_endpoint": api_base,
"azure_deployment": model,
"http_client": litellm.aclient_session,
"max_retries": max_retries,
"timeout": timeout,
}
azure_client_params = select_azure_base_url_or_endpoint(
azure_client_params=azure_client_params
)
if api_key is not None:
azure_client_params["api_key"] = api_key
elif azure_ad_token is not None:
if azure_ad_token.startswith("oidc/"):
azure_ad_token = get_azure_ad_token_from_oidc(azure_ad_token)
azure_client_params["azure_ad_token"] = azure_ad_token
elif azure_ad_token_provider is not None:
azure_client_params["azure_ad_token_provider"] = azure_ad_token_provider
if client is None or dynamic_params: if client is None or dynamic_params:
azure_client = AsyncAzureOpenAI(**azure_client_params) azure_client = AsyncAzureOpenAI(**azure_client_params)
else: else:

44
litellm/router.py
View file

@ -5353,36 +5353,12 @@ class Router:
client = self.cache.get_cache( client = self.cache.get_cache(
key=cache_key, local_only=True, parent_otel_span=parent_otel_span key=cache_key, local_only=True, parent_otel_span=parent_otel_span
) )
if client is None:
"""
Re-initialize the client
"""
InitalizeOpenAISDKClient.set_client(
litellm_router_instance=self, model=deployment
)
client = self.cache.get_cache(
key=cache_key,
local_only=True,
parent_otel_span=parent_otel_span,
)
return client return client
else: else:
cache_key = f"{model_id}_async_client" cache_key = f"{model_id}_async_client"
client = self.cache.get_cache( client = self.cache.get_cache(
key=cache_key, local_only=True, parent_otel_span=parent_otel_span key=cache_key, local_only=True, parent_otel_span=parent_otel_span
) )
# if client is None:
# """
# Re-initialize the client
# """
# InitalizeOpenAISDKClient.set_client(
# litellm_router_instance=self, model=deployment
# )
# client = self.cache.get_cache(
# key=cache_key,
# local_only=True,
# parent_otel_span=parent_otel_span,
# )
return client return client
else: else:
if kwargs.get("stream") is True: if kwargs.get("stream") is True:
@ -5390,32 +5366,12 @@ class Router:
client = self.cache.get_cache( client = self.cache.get_cache(
key=cache_key, parent_otel_span=parent_otel_span key=cache_key, parent_otel_span=parent_otel_span
) )
if client is None:
"""
Re-initialize the client
"""
InitalizeOpenAISDKClient.set_client(
litellm_router_instance=self, model=deployment
)
client = self.cache.get_cache(
key=cache_key, parent_otel_span=parent_otel_span
)
return client return client
else: else:
cache_key = f"{model_id}_client" cache_key = f"{model_id}_client"
client = self.cache.get_cache( client = self.cache.get_cache(
key=cache_key, parent_otel_span=parent_otel_span key=cache_key, parent_otel_span=parent_otel_span
) )
if client is None:
"""
Re-initialize the client
"""
InitalizeOpenAISDKClient.set_client(
litellm_router_instance=self, model=deployment
)
client = self.cache.get_cache(
key=cache_key, parent_otel_span=parent_otel_span
)
return client return client
def _pre_call_checks( # noqa: PLR0915 def _pre_call_checks( # noqa: PLR0915