Merge pull request #9047 from BerriAI/litellm_dev_03_06_2025_p4

feat(handle_jwt.py): support multiple jwt url's
This commit is contained in:
Krish Dholakia 2025-03-10 22:37:35 -07:00 committed by GitHub
commit c93a5e2301
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 71 additions and 36 deletions

View file

@ -345,32 +345,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