From 381347082dbde57d7b60a7fdb1ff8b74b6746c42 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 8 Jul 2024 16:39:05 -0700 Subject: [PATCH] fix(azure.py): improve error handling for azure image gen responses --- litellm/llms/azure.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index f905d17a3c..a2928cf208 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -1108,7 +1108,10 @@ class AzureChatCompletion(BaseLLM): "api-key": api_key, }, ) - operation_location_url = response.headers["operation-location"] + if "operation-location" in response.headers: + operation_location_url = response.headers["operation-location"] + else: + raise AzureOpenAIError(status_code=500, message=response.text) response = await async_handler.get( url=operation_location_url, headers={ @@ -1220,7 +1223,10 @@ class AzureChatCompletion(BaseLLM): "api-key": api_key, }, ) - operation_location_url = response.headers["operation-location"] + if "operation-location" in response.headers: + operation_location_url = response.headers["operation-location"] + else: + raise AzureOpenAIError(status_code=500, message=response.text) response = sync_handler.get( url=operation_location_url, headers={