diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 0371fc6c66..7a7a38b79d 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -111,14 +111,47 @@ class AzureChatCompletion(BaseLLM): exception_mapping_worked = False try: - if (model is None or messages is None) and client is None: + if model is None or messages is None: raise AzureOpenAIError(status_code=422, message=f"Missing model or messages") - - data = { - "model": model, - "messages": messages, - **optional_params - } + + max_retries = optional_params.pop("max_retries", 2) + + ### CHECK IF CLOUDFLARE AI GATEWAY ### + ### if so - set the model as part of the base url + if "gateway.ai.cloudflare.com" in api_base and client is None: + ## build base url - assume api base includes resource name + if not api_base.endswith("/"): + api_base += "/" + api_base += f"{model}" + + azure_client_params = { + "api_version": api_version, + "base_url": f"{api_base}", + "http_client": litellm.client_session, + "max_retries": max_retries, + "timeout": timeout + } + 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 acompletion is True: + client = AsyncAzureOpenAI(**azure_client_params) + else: + client = AzureOpenAI(**azure_client_params) + + data = { + "model": None, + "messages": messages, + **optional_params + } + else: + data = { + "model": model, # type: ignore + "messages": messages, + **optional_params + } ## LOGGING logging_obj.pre_call( input=messages, @@ -133,7 +166,7 @@ class AzureChatCompletion(BaseLLM): "complete_input_dict": data, }, ) - max_retries = data.pop("max_retries", 2) + if acompletion is True: if optional_params.get("stream", False): return self.async_streaming(logging_obj=logging_obj, api_base=api_base, data=data, model=model, api_key=api_key, api_version=api_version, azure_ad_token=azure_ad_token, timeout=timeout, client=client) diff --git a/litellm/main.py b/litellm/main.py index fa5a8c868a..850d256723 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -494,31 +494,6 @@ def completion( if k not in optional_params: # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in optional_params[k] = v - ### CHECK IF CLOUDFLARE AI GATEWAY ### - ### if so - set the model as part of the base url - if "gateway.ai.cloudflare.com" in api_base and client is None: - ## build base url - assume api base includes resource name - if not api_base.endswith("/"): - api_base += "/" - api_base += f"{model}" - - azure_client_params = { - "api_version": api_version, - "base_url": f"{api_base}", - "http_client": litellm.client_session, - "max_retries": max_retries, - "timeout": timeout - } - 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 acompletion is True: - client = openai.AsyncAzureOpenAI(**azure_client_params) - else: - client = openai.AzureOpenAI(**azure_client_params) - model = None - ## COMPLETION CALL response = azure_chat_completions.completion( model=model, diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 8412d0f192..dcc7719907 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1334,7 +1334,7 @@ def test_azure_cloudflare_api(): response = completion(model="azure/gpt-turbo", messages=messages, base_url=os.getenv("CLOUDFLARE_AZURE_BASE_URL"), api_key=os.getenv("AZURE_FRANCE_API_KEY")) print(f"response: {response}") except Exception as e: - print(e) + traceback.print_exc() pass test_azure_cloudflare_api()