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

feat(handle_jwt.py): support multiple jwt url's

Browse source
This commit is contained in:
Krrish Dholakia 2025-03-06 23:05:54 -08:00
parent 5b2eb1f6bb
commit 5280a914cd
2 changed files with 70 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

50
litellm/proxy/auth/handle_jwt.py
View file

@ -344,32 +344,38 @@ class JWTHandler:
if keys_url is None:
raise Exception("Missing JWT Public Key URL from environment.")
cached_keys = await self.user_api_key_cache.async_get_cache(
"litellm_jwt_auth_keys"
)
if cached_keys is None:
response = await self.http_handler.get(keys_url)
keys_url_list = [url.strip() for url in keys_url.split(",")]
response_json = response.json()
if "keys" in response_json:
keys: JWKKeyValue = response.json()["keys"]
for key_url in keys_url_list:
cache_key = f"litellm_jwt_auth_keys_{key_url}"
cached_keys = await self.user_api_key_cache.async_get_cache(cache_key)
if cached_keys is None:
response = await self.http_handler.get(key_url)
response_json = response.json()
if "keys" in response_json:
keys: JWKKeyValue = response.json()["keys"]
else:
keys = response_json
await self.user_api_key_cache.async_set_cache(
key=cache_key,
value=keys,
ttl=self.litellm_jwtauth.public_key_ttl, # cache for 10 mins
)
else:
keys = response_json
keys = cached_keys
await self.user_api_key_cache.async_set_cache(
key="litellm_jwt_auth_keys",
value=keys,
ttl=self.litellm_jwtauth.public_key_ttl, # cache for 10 mins
)
else:
keys = cached_keys
public_key = self.parse_keys(keys=keys, kid=kid)
if public_key is not None:
return cast(dict, public_key)
public_key = self.parse_keys(keys=keys, kid=kid)
if public_key is None:
raise Exception(
f"No matching public key found. kid={kid}, keys_url={keys_url}, cached_keys={cached_keys}, len(keys)={len(keys)}"
)
return cast(dict, public_key)
raise Exception(
f"No matching public key found. keys={keys_url_list}, kid={kid}"
)
def parse_keys(self, keys: JWKKeyValue, kid: Optional[str]) -> Optional[JWTKeyItem]:
public_key: Optional[JWTKeyItem] = None