From 145cd483b924541cbaadb0706b4556e316220eb5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 11 Mar 2025 18:48:12 -0700 Subject: [PATCH] fix(handler.py): same as last commit --- litellm/llms/azure/completion/handler.py | 81 +++++------------------- 1 file changed, 16 insertions(+), 65 deletions(-) diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index fafa5665bb..186bc067a9 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -8,7 +8,7 @@ from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResp from ...base import BaseLLM from ...openai.completion.transformation import OpenAITextCompletionConfig -from ..common_utils import AzureOpenAIError +from ..common_utils import AzureOpenAIError, BaseAzureLLM openai_text_completion_config = OpenAITextCompletionConfig() @@ -25,7 +25,7 @@ def select_azure_base_url_or_endpoint(azure_client_params: dict): return azure_client_params -class AzureTextCompletion(BaseLLM): +class AzureTextCompletion(BaseAzureLLM): def __init__(self) -> None: super().__init__() @@ -60,7 +60,6 @@ class AzureTextCompletion(BaseLLM): headers: Optional[dict] = None, client=None, ): - super().completion() try: if model is None or messages is None: raise AzureOpenAIError( @@ -72,6 +71,14 @@ class AzureTextCompletion(BaseLLM): messages=messages, model=model, custom_llm_provider="azure_text" ) + azure_client_params = self.initialize_azure_sdk_client( + litellm_params=litellm_params or {}, + api_key=api_key, + model_name=model, + api_version=api_version, + api_base=api_base, + ) + ### CHECK IF CLOUDFLARE AI GATEWAY ### ### if so - set the model as part of the base url if "gateway.ai.cloudflare.com" in api_base: @@ -118,6 +125,7 @@ class AzureTextCompletion(BaseLLM): azure_ad_token=azure_ad_token, timeout=timeout, client=client, + azure_client_params=azure_client_params, ) else: return self.acompletion( @@ -132,6 +140,7 @@ class AzureTextCompletion(BaseLLM): client=client, logging_obj=logging_obj, max_retries=max_retries, + azure_client_params=azure_client_params, ) elif "stream" in optional_params and optional_params["stream"] is True: return self.streaming( @@ -144,6 +153,7 @@ class AzureTextCompletion(BaseLLM): azure_ad_token=azure_ad_token, timeout=timeout, client=client, + azure_client_params=azure_client_params, ) else: ## LOGGING @@ -165,22 +175,6 @@ class AzureTextCompletion(BaseLLM): status_code=422, message="max retries must be an int" ) # 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_ad_token_provider": azure_ad_token_provider, - } - 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: - azure_client_params["azure_ad_token"] = azure_ad_token if client is None: azure_client = AzureOpenAI(**azure_client_params) else: @@ -240,26 +234,11 @@ class AzureTextCompletion(BaseLLM): max_retries: int, azure_ad_token: Optional[str] = None, client=None, # this is the AsyncAzureOpenAI + azure_client_params: dict = {}, ): response = None try: # 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: - azure_client_params["azure_ad_token"] = azure_ad_token - # setting Azure client if client is None: azure_client = AsyncAzureOpenAI(**azure_client_params) @@ -312,6 +291,7 @@ class AzureTextCompletion(BaseLLM): timeout: Any, azure_ad_token: Optional[str] = None, client=None, + azure_client_params: dict = {}, ): max_retries = data.pop("max_retries", 2) if not isinstance(max_retries, int): @@ -319,21 +299,6 @@ class AzureTextCompletion(BaseLLM): status_code=422, message="max retries must be an int" ) # 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: - azure_client_params["azure_ad_token"] = azure_ad_token if client is None: azure_client = AzureOpenAI(**azure_client_params) else: @@ -375,24 +340,10 @@ class AzureTextCompletion(BaseLLM): timeout: Any, azure_ad_token: Optional[str] = None, client=None, + azure_client_params: dict = {}, ): try: # init AzureOpenAI Client - azure_client_params = { - "api_version": api_version, - "azure_endpoint": api_base, - "azure_deployment": model, - "http_client": litellm.client_session, - "max_retries": data.pop("max_retries", 2), - "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: - azure_client_params["azure_ad_token"] = azure_ad_token if client is None: azure_client = AsyncAzureOpenAI(**azure_client_params) else: