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

fix(handle_jwt.py): enable user to set jwt admin scope string in config

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-25 11:34:49 -07:00
parent d2d6f6d0aa
commit d3e7376e10
2 changed files with 11 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/_types.py
View file

@ -14,11 +14,6 @@ def hash_token(token: str):
return hashed_token return hashed_token
class LiteLLMProxyRoles(enum.Enum):
PROXY_ADMIN = "litellm_proxy_admin"
USER = "litellm_user"
class LiteLLMBase(BaseModel): class LiteLLMBase(BaseModel):
""" """
Implements default functions, all pydantic objects should have. Implements default functions, all pydantic objects should have.
@ -42,6 +37,11 @@ class LiteLLMBase(BaseModel):
protected_namespaces = () protected_namespaces = ()
class LiteLLMProxyRoles(LiteLLMBase):
PROXY_ADMIN: str = "litellm_proxy_admin"
PROXY_USER: str = "litellm_user"
class LiteLLMPromptInjectionParams(LiteLLMBase): class LiteLLMPromptInjectionParams(LiteLLMBase):
heuristics_check: bool = False heuristics_check: bool = False
vector_db_check: bool = False vector_db_check: bool = False

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

@ -67,17 +67,21 @@ class JWTHandler:
self.http_handler = HTTPHandler() self.http_handler = HTTPHandler()
def update_environment( def update_environment(
self, prisma_client: Optional[PrismaClient], user_api_key_cache: DualCache self,
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
litellm_proxy_roles: LiteLLMProxyRoles,
) -> None: ) -> None:
self.prisma_client = prisma_client self.prisma_client = prisma_client
self.user_api_key_cache = user_api_key_cache self.user_api_key_cache = user_api_key_cache
self.litellm_proxy_roles = litellm_proxy_roles
def is_jwt(self, token: str): def is_jwt(self, token: str):
parts = token.split(".") parts = token.split(".")
return len(parts) == 3 return len(parts) == 3
def is_admin(self, scopes: list) -> bool: def is_admin(self, scopes: list) -> bool:
if LiteLLMProxyRoles.PROXY_ADMIN.value in scopes: if self.litellm_proxy_roles.PROXY_ADMIN in scopes:
return True return True
return False return False