From 5c0d4a4162c74b93a1fc023f58bea5221dcec8d4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 8 Jul 2024 12:47:53 -0700 Subject: [PATCH] use ProxyErrorTypes, --- .../health_endpoints/_health_endpoints.py | 12 ++- .../management_endpoints/team_endpoints.py | 100 +++++++++--------- 2 files changed, 60 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 68778e177..29ad6bae7 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -10,7 +10,13 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import CallInfo, ProxyException, UserAPIKeyAuth, WebhookEvent +from litellm.proxy._types import ( + CallInfo, + ProxyErrorTypes, + ProxyException, + UserAPIKeyAuth, + WebhookEvent, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.health_check import perform_health_check @@ -239,7 +245,7 @@ async def health_services_endpoint( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type="auth_error", + type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) @@ -247,7 +253,7 @@ async def health_services_endpoint( raise e raise ProxyException( message="Authentication Error, " + str(e), - type="auth_error", + type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 50244ee23..6be8310e8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1,32 +1,34 @@ -from typing import Optional, List -import fastapi -from fastapi import Depends, Request, APIRouter, Header, status -from fastapi import HTTPException +import asyncio import copy import json import uuid -import litellm -import asyncio from datetime import datetime, timedelta, timezone +from typing import List, Optional + +import fastapi +from fastapi import APIRouter, Depends, Header, HTTPException, Request, status + +import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy._types import ( - UserAPIKeyAuth, - LiteLLM_TeamTable, - LiteLLM_ModelTable, - LitellmUserRoles, - NewTeamRequest, - TeamMemberAddRequest, - UpdateTeamRequest, BlockTeamRequest, - DeleteTeamRequest, - Member, - LitellmTableNames, - LiteLLM_AuditLogs, - TeamMemberDeleteRequest, - ProxyException, CommonProxyErrors, + DeleteTeamRequest, + LiteLLM_AuditLogs, + LiteLLM_ModelTable, + LiteLLM_TeamTable, + LitellmTableNames, + LitellmUserRoles, + Member, + NewTeamRequest, + ProxyErrorTypes, + ProxyException, + TeamMemberAddRequest, + TeamMemberDeleteRequest, + UpdateTeamRequest, + UserAPIKeyAuth, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, @@ -109,10 +111,10 @@ async def new_team( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -326,10 +328,10 @@ async def update_team( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -420,10 +422,10 @@ async def team_member_add( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -525,10 +527,10 @@ async def team_member_delete( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -639,10 +641,10 @@ async def delete_team( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -725,10 +727,10 @@ async def team_info( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) try: @@ -783,7 +785,7 @@ async def team_info( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type="auth_error", + type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), ) @@ -791,7 +793,7 @@ async def team_info( raise e raise ProxyException( message="Authentication Error, " + str(e), - type="auth_error", + type=ProxyErrorTypes.auth_error, param=getattr(e, "param", "None"), code=status.HTTP_400_BAD_REQUEST, ) @@ -810,10 +812,10 @@ async def block_team( Blocks all calls from keys with this team id. """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -839,10 +841,10 @@ async def unblock_team( Blocks all calls from keys with this team id. """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if prisma_client is None: @@ -872,10 +874,10 @@ async def list_team( ``` """ from litellm.proxy.proxy_server import ( - prisma_client, - litellm_proxy_admin_name, - create_audit_log_for_update, _duration_in_seconds, + create_audit_log_for_update, + litellm_proxy_admin_name, + prisma_client, ) if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: