BerriAI / litellm

Python SDK, Proxy Server (LLM Gateway) to call 100+ LLM APIs in OpenAI format - [Bedrock, Azure, OpenAI, VertexAI, Cohere, Anthropic, Sagemaker, HuggingFace, Replicate, Groq]
https://docs.litellm.ai/docs/
Other
13.91k stars 1.64k forks source link

[Bug]: Error when trying to view Usage as an internal_user #6238

Closed ncecere closed 3 weeks ago

ncecere commented 1 month ago

What happened?

In the UI I was trying to view my spend usage and I received the following error.

image

Now I'm able to see some information in the Usage section.

image

I'm an enterprise customer and we are using OIDC for auth.

I'm wondering if this issue is related to the when I logged in for the first time a Key was automatically generated for my user.

image

we are using a custom_sso_handler durring login to limit the models users can see.

import os
from fastapi_sso.sso.base import OpenID
from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues, Member, UserAPIKeyAuth, NewUserRequest, UpdateUserRequest
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user, user_update
from litellm.proxy.management_helpers.utils import add_new_member
from litellm import default_internal_user_params
from litellm.proxy.utils import PrismaClient
from litellm.proxy.proxy_server import proxy_logging_obj
import datetime

# Initialize prisma_client with required arguments
database_url = os.getenv("DATABASE_URL")
sso_key = os.getenv("SSO_KEY")
if not database_url:
    raise EnvironmentError("DATABASE_URL environment variable is not set")

prisma_client = PrismaClient(
    database_url=database_url,
    proxy_logging_obj=proxy_logging_obj
)

# async def get_team_id_by_alias(team_alias: str) -> str:
#     """
#     Retrieve the team ID based on the team alias.
#     """
#     if prisma_client is None:
#         raise Exception("Not connected to DB!")

#     team_info = await prisma_client.db.litellm_teamtable.find_first(
#         where={"team_alias": team_alias}
#     )

#     if team_info is None:
#         raise ValueError(f"Team alias '{team_alias}' does not exist in the database.")

#     return team_info.team_id

async def get_user_info(user_id: str):
    """
    Retrieve user information from the LiteLLM DB.
    """
    if prisma_client is None:
        raise Exception("Not connected to DB!")

    user_info = await prisma_client.db.litellm_usertable.find_unique(
        where={"user_id": user_id}
    )

    return user_info

async def is_user_in_team(user_id: str, team_id: str) -> bool:
    """
    Check if the user is already a member of the team.
    """
    if prisma_client is None:
        raise Exception("Not connected to DB!")

    membership = await prisma_client.db.litellm_teammembership.find_first(
        where={"user_id": user_id, "team_id": team_id}
    )

    return membership is not None

async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues:
    try:
        print("inside custom sso handler")  # noqa
        print(f"userIDPInfo: {userIDPInfo}")  # noqa

        if userIDPInfo.id is None:
            raise ValueError(
                f"No ID found for user. userIDPInfo.id is None {userIDPInfo}"
            )

        # Connect to the Prisma client
        await prisma_client.connect()

        # Extract model_list and team_alias from default_internal_user_params
        model_list = default_internal_user_params["models"]
        # team_alias = default_internal_user_params["team_alias"]

        # Check if user exists in LiteLLM DB
        _user_info = await get_user_info(user_id=userIDPInfo.id)

        # team_id = await get_team_id_by_alias(team_alias)

        if _user_info is None:
            # User does not exist, create a new user
            new_user_data = NewUserRequest(
                user_id=userIDPInfo.id,
                user_email=userIDPInfo.email,
                user_role=LitellmUserRoles.INTERNAL_USER.value,
                max_budget=default_internal_user_params["max_budget"],
                budget_duration=default_internal_user_params["budget_duration"],
                models=model_list,
            )
            await new_user(new_user_data)
        else:
            # User exists, update user information
            update_user_data = UpdateUserRequest(
                user_id=userIDPInfo.id,
                max_budget=default_internal_user_params["max_budget"],
                budget_duration=default_internal_user_params["budget_duration"],
                models=model_list,
            )
            await user_update(update_user_data)

        # Retrieve the master key from environment variable
        master_key = sso_key

        # Create a minimal user_api_key_dict
        user_api_key_dict = UserAPIKeyAuth(
            user_id=userIDPInfo.id,
            user_role=LitellmUserRoles.INTERNAL_USER.value,
            api_key=master_key  # Set the API key from environment variable
        )

        # Check if the user is already a member of the team
        # if not await is_user_in_team(userIDPInfo.id, team_id):
        #     # Add user to team using add_new_member
        #     new_member = Member(
        #         user_id=userIDPInfo.id,
        #         user_email=userIDPInfo.email,
        #         role="user"  # Map internal_user to user
        #     )
        #     await add_new_member(
        #         new_member=new_member,
        #         max_budget_in_team=default_internal_user_params["max_budget"],
        #         prisma_client=prisma_client,
        #         team_id=team_id,
        #         user_api_key_dict=user_api_key_dict,
        #         litellm_proxy_admin_name="halbert"  # Replace with actual admin name if available
        #     )

        return SSOUserDefinedValues(
            models=model_list,                              # models user has access to
            user_id=userIDPInfo.id,                         # user id to use in the LiteLLM DB
            user_email=userIDPInfo.email,                   # user email to use in the LiteLLM DB
            user_role=LitellmUserRoles.INTERNAL_USER.value, # role to use for the user
            max_budget=default_internal_user_params["max_budget"],  # Max budget for this UI login Session
            budget_duration=default_internal_user_params["budget_duration"],  # Duration of the budget for this UI login Session, 1d, 2d, 30d ...
        )
    except Exception as e:
        raise Exception("Failed custom auth") from e
    finally:
        # Disconnect the Prisma client
        await prisma_client.disconnect()

Relevant log output

Let me know if there is any additional information that I need to provide, I'll make sure to get it to you right away

Twitter / LinkedIn details

No response

ncecere commented 1 month ago

I get the following error when I try to view usage

{"error":{"message":"Authentication Error, Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route=/global/spend/logs. Your role=internal_user. Your user_id=nicholas.cecere@ufl.edu","type":"auth_error","param":"None","code":"401"}}

also after getting this error, if I go to virtual keys the budget for the default team will show No Limit even if I have a default_internal_user_budget set

krrishdholakia commented 4 weeks ago

able to repro