feat(health_check.py): more detailed health check calls

litellm/llms/azure.py

@@ -700,3 +700,72 @@ class AzureChatCompletion(BaseLLM):
                 import traceback
 
                 raise AzureOpenAIError(status_code=500, message=traceback.format_exc())
+
+    async def ahealth_check(
+        self,
+        model: Optional[str],
+        api_key: str,
+        api_base: str,
+        api_version: str,
+        timeout: float,
+        mode: str,
+        messages: Optional[list] = None,
+        input: Optional[list] = None,
+        prompt: Optional[str] = None,
+    ):
+        client_session = litellm.aclient_session or httpx.AsyncClient(
+            transport=AsyncCustomHTTPTransport(),  # handle dall-e-2 calls
+        )
+        client = AsyncAzureOpenAI(
+            api_version=api_version,
+            azure_endpoint=api_base,
+            api_key=api_key,
+            timeout=timeout,
+            http_client=client_session,
+        )
+
+        if model is None and mode != "image_generation":
+            raise Exception("model is not set")
+
+        completion = None
+
+        if mode == "completion":
+            if messages is None:
+                raise Exception("messages is not set")
+            completion = await client.chat.completions.with_raw_response.create(
+                model=model,  # type: ignore
+                messages=messages,  # type: ignore
+            )
+        elif mode == "embedding":
+            if input is None:
+                raise Exception("input is not set")
+            completion = await client.embeddings.with_raw_response.create(
+                model=model,  # type: ignore
+                input=input,  # type: ignore
+            )
+        elif mode == "image_generation":
+            if prompt is None:
+                raise Exception("prompt is not set")
+            completion = await client.images.with_raw_response.generate(
+                model=model,  # type: ignore
+                prompt=prompt,  # type: ignore
+            )
+        else:
+            raise Exception("mode not set")
+        response = {}
+
+        if completion is None or not hasattr(completion, "headers"):
+            raise Exception("invalid completion response")
+
+        if (
+            completion.headers.get("x-ratelimit-remaining-requests", None) is not None
+        ):  # not provided for dall-e requests
+            response["x-ratelimit-remaining-requests"] = completion.headers[
+                "x-ratelimit-remaining-requests"
+            ]
+
+        if completion.headers.get("x-ratelimit-remaining-tokens", None) is not None:
+            response["x-ratelimit-remaining-tokens"] = completion.headers[
+                "x-ratelimit-remaining-tokens"
+            ]
+        return response