Add support for upserting users automatically to a default team based on JWT key #3717
Conversation
…C provider New configuration option team_name_default: "TEAMA" # Will add the user to this team by default user_email_jwt_field: "email" # will associate email to user when they are created in the db
|
The latest updates on your projects. Learn more about Vercel for Git ↗︎
|
| @@ -227,6 +227,14 @@ class LiteLLM_JWTAuth(LiteLLMBase): | |||
| user_id_upsert: bool = Field( | |||
| default=False, description="If user doesn't exist, upsert them into the db." | |||
| ) | |||
| user_email_jwt_field: Optional[str] = Field( | |||
| default="email", | |||
There was a problem hiding this comment.
@Collen-Roller why use a default str instead of leaving it as None?
There was a problem hiding this comment.
Because most OIDC providers use the email and this would automatically associate it the user that gets created. It's an entry in the DB so wouldn't it make sense to have a default so the user doesn't have to configure it? Manual override option with user_email_jwt_field?
litellm/proxy/auth/auth_checks.py
Outdated
| where={"team_alias": team_name_default} | ||
| ) | ||
|
|
||
| print("RESPONSE AFTER FINDING TEAM FROM ALIAS") |
There was a problem hiding this comment.
can we cleanup the print statements - it'll fail our linting
There was a problem hiding this comment.
Yes I'll clean this up. Done
| user_to_add["teams"] = [response.data.team_id] | ||
|
|
||
| except Exception as e: # if user not in db | ||
| # Create the team id automatically |
There was a problem hiding this comment.
the team id upsert logic should not be happening inside 'get_user_object'. this is a generic function, meant to be used for both key + jwt_auth
and to just get the user object if it exists
There was a problem hiding this comment.
if team_id upserts need to happen, they should happen within the 'get_team_object' function
There was a problem hiding this comment.
Tried to modify this and couldn't get this working. Logic I believe should stay in the user based on workflow. If a user with a key tries to hit LLM for the first time, the team should be automatically created. Once the team is created, users will be automatically joined to the team
| user_to_add["teams"] = [response.data.team_id] | ||
|
|
||
| except Exception as e: # if user not in db | ||
| # Create the team id automatically |
There was a problem hiding this comment.
if team_id upserts need to happen, they should happen within the 'get_team_object' function
| else: | ||
| user_email = None | ||
| except KeyError: | ||
| user_email = default_value |
There was a problem hiding this comment.
why are we returning a default value for user_email?
will all users use the same email?
There was a problem hiding this comment.
Using the same logic for user_id - default_value is None
litellm/proxy/proxy_server.py
Outdated
| @@ -352,31 +352,6 @@ def _get_pydantic_json_dict(pydantic_obj: BaseModel) -> dict: | |||
| return pydantic_obj.dict() | |||
|
|
|||
|
|
|||
| def get_custom_headers( | |||
There was a problem hiding this comment.
your PR seems to overwrite updates to proxy_server.py, can you rebase, to avoid this/
|
finished reviewing PR. Updated repo |
Added features to support new JWT options
Fixes
https://<YOUR_OIDC_PROVIDER>/auth/realms/dev/protocol/openid-connect/certs
Code fix: Fixed JWT public key finding #3648
Test add: Adding multiple public keys test #3649
Type
🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
Changes
user_id_upsert1.1. Added
team_name_default- This allows users to specific a name for the team the user would like to be added to by default. This takes away the need to specific default team ID which NEEDS to exist beforehand. Now, if the team name doesn't exist, a new one gets created.1.2. Added
user_email_jwt_field- This allows the user to specify what attribute from the JWT will be used. If the attribute exists in the JWT, it will upsert the users email address into their record for tracking.[REQUIRED] Testing - Attach a screenshot of any new tests passing locall
export DATABASE_URL=postgresql://:@:/<DATABASE_NAME>
export JWT_PUBLIC_KEY_URL=this_doesnt_matter # ...because the script makes its own pub/priv keys on the fly
prisma generate
prisma db push