Merge pull request #5209 from BerriAI/litellm_azure_ai_openai_support

fix(utils.py): support calling openai models via `azure_ai/`

litellm/main.py

@@ -4898,7 +4898,6 @@ async def ahealth_check(
         verbose_logger.error(
             "litellm.ahealth_check(): Exception occured - {}".format(str(e))
         )
-        verbose_logger.debug(traceback.format_exc())
         stack_trace = traceback.format_exc()
         if isinstance(stack_trace, str):
             stack_trace = stack_trace[:1000]
@@ -4907,7 +4906,12 @@ async def ahealth_check(
                 "error": "Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models"
             }
 
-        error_to_return = str(e) + " stack trace: " + stack_trace
+        error_to_return = (
+            str(e)
+            + "\nHave you set 'mode' - https://docs.litellm.ai/docs/proxy/health#embedding-models"
+            + "\nstack trace: "
+            + stack_trace
+        )
         return {"error": error_to_return}
 
 

litellm/proxy/_new_secret_config.yaml

@@ -3,5 +3,4 @@ model_list:
     litellm_params:
       model: "gpt-4"
     model_info:
-      my_custom_key: "my_custom_value"
-
+      my_custom_key: "my_custom_value"

litellm/tests/test_completion.py

@@ -190,6 +190,31 @@ def test_completion_azure_command_r():
         pytest.fail(f"Error occurred: {e}")
 
 
+@pytest.mark.parametrize(
+    "api_base",
+    [
+        "https://litellm8397336933.openai.azure.com",
+        "https://litellm8397336933.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2023-03-15-preview",
+    ],
+)
+def test_completion_azure_ai_gpt_4o(api_base):
+    try:
+        litellm.set_verbose = True
+
+        response = completion(
+            model="azure_ai/gpt-4o",
+            api_base=api_base,
+            api_key=os.getenv("AZURE_AI_OPENAI_KEY"),
+            messages=[{"role": "user", "content": "What is the meaning of life?"}],
+        )
+
+        print(response)
+    except litellm.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 @pytest.mark.parametrize("sync_mode", [True, False])
 @pytest.mark.asyncio
 async def test_completion_databricks(sync_mode):

litellm/utils.py

@@ -4479,7 +4479,22 @@ def _is_non_openai_azure_model(model: str) -> bool:
             or f"mistral/{model_name}" in litellm.mistral_chat_models
         ):
             return True
-    except:
+    except Exception:
+        return False
+    return False
+
+
+def _is_azure_openai_model(model: str) -> bool:
+    try:
+        if "/" in model:
+            model = model.split("/", 1)[1]
+        if (
+            model in litellm.open_ai_chat_completion_models
+            or model in litellm.open_ai_text_completion_models
+            or model in litellm.open_ai_embedding_models
+        ):
+            return True
+    except Exception:
         return False
     return False
 
@@ -4613,6 +4628,14 @@ def get_llm_provider(
             elif custom_llm_provider == "azure_ai":
                 api_base = api_base or get_secret("AZURE_AI_API_BASE")  # type: ignore
                 dynamic_api_key = api_key or get_secret("AZURE_AI_API_KEY")
+
+                if _is_azure_openai_model(model=model):
+                    verbose_logger.debug(
+                        "Model={} is Azure OpenAI model. Setting custom_llm_provider='azure'.".format(
+                            model
+                        )
+                    )
+                    custom_llm_provider = "azure"
             elif custom_llm_provider == "github":
                 api_base = api_base or get_secret("GITHUB_API_BASE") or "https://models.inference.ai.azure.com"  # type: ignore
                 dynamic_api_key = api_key or get_secret("GITHUB_API_KEY")