phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(azure.py): Add Azure AD cred caching for OIDC flow.

Browse source
This commit is contained in:
David Manouchehri 2024-05-27 17:35:37 +00:00
parent a31fa5fbc8
commit d950373768
No known key found for this signature in database
1 changed files with 29 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

33
litellm/llms/azure.py
View file

@ -36,6 +36,9 @@ from ..types.llms.openai import (
AsyncAssistantStreamManager, AsyncAssistantStreamManager,
AssistantStreamManager, AssistantStreamManager,
) )
from litellm.caching import DualCache
azure_ad_cache = DualCache()
class AzureOpenAIError(Exception): class AzureOpenAIError(Exception):
@ -326,6 +329,17 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str):
message="OIDC token could not be retrieved from secret manager.", message="OIDC token could not be retrieved from secret manager.",
) )
azure_ad_token_cache_key = json.dumps({
"azure_client_id": azure_client_id,
"azure_tenant_id": azure_tenant_id,
"azure_authority_host": azure_authority_host,
"oidc_token": oidc_token,
})
azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key)
if azure_ad_token_access_token is not None:
return azure_ad_token_access_token
req_token = httpx.post( req_token = httpx.post(
f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
data={ data={
@ -343,12 +357,23 @@ def get_azure_ad_token_from_oidc(azure_ad_token: str):
message=req_token.text, message=req_token.text,
) )
possible_azure_ad_token = req_token.json().get("access_token", None) azure_ad_token_json = req_token.json()
azure_ad_token_access_token = azure_ad_token_json.get("access_token", None)
azure_ad_token_expires_in = azure_ad_token_json.get("expires_in", None)
if possible_azure_ad_token is None: if azure_ad_token_access_token is None:
raise AzureOpenAIError(status_code=422, message="Azure AD Token not returned") raise AzureOpenAIError(
status_code=422, message="Azure AD Token access_token not returned"
)
return possible_azure_ad_token if azure_ad_token_expires_in is None:
raise AzureOpenAIError(
status_code=422, message="Azure AD Token expires_in not returned"
)
azure_ad_cache.set_cache(key=azure_ad_token_cache_key, value=azure_ad_token_access_token, ttl=azure_ad_token_expires_in)
return azure_ad_token_access_token
class AzureChatCompletion(BaseLLM): class AzureChatCompletion(BaseLLM):