phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat) add azure_ad_token

Browse source
This commit is contained in:
ishaan-jaff 2023-11-14 12:09:50 -08:00
parent ceaf944cd4
commit ed56f280d9
1 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/llms/azure.py
View file

@ -78,12 +78,14 @@ class AzureChatCompletion(BaseLLM):
super().__init__() super().__init__()
self._client_session = self.create_client_session() self._client_session = self.create_client_session()
def validate_environment(self, api_key): def validate_environment(self, api_key, azure_ad_token):
headers = { headers = {
"content-type": "application/json", "content-type": "application/json",
} }
if api_key: if api_key is not None:
headers["api-key"] = api_key headers["api-key"] = api_key
if azure_ad_token is not None:
headers["Authorization"] = f"Bearer {azure_ad_token}"
return headers return headers
def completion(self, def completion(self,
@ -94,6 +96,7 @@ class AzureChatCompletion(BaseLLM):
api_base: str, api_base: str,
api_version: str, api_version: str,
api_type: str, api_type: str,
azure_ad_token: str,
print_verbose: Callable, print_verbose: Callable,
logging_obj, logging_obj,
optional_params, optional_params,
@ -105,7 +108,7 @@ class AzureChatCompletion(BaseLLM):
exception_mapping_worked = False exception_mapping_worked = False
try: try:
if headers is None: if headers is None:
headers = self.validate_environment(api_key=api_key) headers = self.validate_environment(api_key=api_key, azure_ad_token=azure_ad_token)
if model is None or messages is None: if model is None or messages is None:
raise AzureOpenAIError(status_code=422, message=f"Missing model or messages") raise AzureOpenAIError(status_code=422, message=f"Missing model or messages")