diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6d2d6e36c..7a6e400c8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -239,6 +239,10 @@ class NewTeamResponse(LiteLLMBase): updated_at: datetime +class TeamRequest(LiteLLMBase): + teams: List[str] + + class KeyManagementSystem(enum.Enum): GOOGLE_KMS = "google_kms" AZURE_KEY_VAULT = "azure_key_vault" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e24fc1a82..ff10b5a17 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2819,7 +2819,8 @@ async def generate_key_fn( Parameters: - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - key_alias: Optional[str] - User defined key alias - - team_id: Optional[str] - The team id of the user + - team_id: Optional[str] - The team id of the key + - user_id: Optional[str] - The user id of the key - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models) - aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models - config: Optional[dict] - any key-specific configs, overrides config in config.yaml @@ -3048,8 +3049,8 @@ async def info_key_fn_v2( Example Curl: ``` curl -X GET "http://0.0.0.0:8000/key/info" \ --H "Authorization: Bearer sk-1234" \ --d {"keys": ["sk-1", "sk-2", "sk-3"]} + -H "Authorization: Bearer sk-1234" \ + -d {"keys": ["sk-1", "sk-2", "sk-3"]} ``` """ global prisma_client @@ -3755,14 +3756,85 @@ async def delete_team(): pass -@router.post( +@router.get( "/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) -async def info_team(): +async def team_info( + team_id: str = fastapi.Query( + default=None, description="Team ID in the request parameters" + ) +): """ get info on team + related keys + + ``` + curl --location 'http://localhost:4000/team/info' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "teams": ["",..] + }' + ``` """ - pass + global prisma_client + try: + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "error": f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + }, + ) + if team_id is None: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Malformed request. No team id passed in."}, + ) + + team_info = await prisma_client.get_data( + team_id=team_id, table_name="team", query_type="find_unique" + ) + ## GET ALL KEYS ## + keys = await prisma_client.get_data( + team_id=team_id, + table_name="key", + query_type="find_all", + expires=datetime.now(), + ) + + if team_info is None: + ## make sure we still return a total spend ## + spend = 0 + for k in keys: + spend += getattr(k, "spend", 0) + team_info = {"spend": spend} + + ## REMOVE HASHED TOKEN INFO before returning ## + for key in keys: + try: + key = key.model_dump() # noqa + except: + # if using pydantic v1 + key = key.dict() + key.pop("token", None) + return {"team_id": team_id, "team_info": team_info, "keys": keys} + + except Exception as e: + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type="auth_error", + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type="auth_error", + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) #### MODEL MANAGEMENT #### diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0108836fc..d82a7231f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -509,8 +509,9 @@ class PrismaClient: token: Optional[Union[str, list]] = None, user_id: Optional[str] = None, user_id_list: Optional[list] = None, + team_id: Optional[str] = None, key_val: Optional[dict] = None, - table_name: Optional[Literal["user", "key", "config", "spend"]] = None, + table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None, query_type: Literal["find_unique", "find_all"] = "find_unique", expires: Optional[datetime] = None, reset_at: Optional[datetime] = None, @@ -545,6 +546,14 @@ class PrismaClient: for r in response: if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() + elif query_type == "find_all" and team_id is not None: + response = await self.db.litellm_verificationtoken.find_many( + where={"team_id": team_id} + ) + if response is not None and len(response) > 0: + for r in response: + if isinstance(r.expires, datetime): + r.expires = r.expires.isoformat() elif ( query_type == "find_all" and expires is not None @@ -657,7 +666,23 @@ class PrismaClient: order={"startTime": "desc"}, ) return response + elif table_name == "team": + if query_type == "find_unique": + response = await self.db.litellm_teamtable.find_unique( + where={"team_id": team_id} # type: ignore + ) + if query_type == "find_all" and team_id is not None: + user_id_values = str(tuple(team_id)) + sql_query = f""" + SELECT * + FROM "LiteLLM_TeamTable" + WHERE "team_id" IN {team_id} + """ + # Execute the raw query + # The asterisk before `team_id` unpacks the list into separate arguments + response = await self.db.query_raw(sql_query) + return response except Exception as e: print_verbose(f"LiteLLM Prisma Client Exception: {e}") import traceback