mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
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:
commit
c93a5e2301
3 changed files with 71 additions and 36 deletions
|
@ -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
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue