forked from phoenix/litellm-mirror
fix(proxy_server.py): rename beta end user blocking endpoints
separating end user management from user management
This commit is contained in:
parent
dd151869a3
commit
844d6828c8
2 changed files with 92 additions and 4 deletions
|
@ -5095,7 +5095,7 @@ async def user_get_requests():
|
||||||
|
|
||||||
|
|
||||||
@router.post(
|
@router.post(
|
||||||
"/user/block",
|
"/end_user/block",
|
||||||
tags=["End User Management"],
|
tags=["End User Management"],
|
||||||
dependencies=[Depends(user_api_key_auth)],
|
dependencies=[Depends(user_api_key_auth)],
|
||||||
)
|
)
|
||||||
|
@ -5133,11 +5133,12 @@ async def block_user(data: BlockUsers):
|
||||||
|
|
||||||
return {"blocked_users": records}
|
return {"blocked_users": records}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
verbose_proxy_logger.error(f"An error occurred - {str(e)}")
|
||||||
raise HTTPException(status_code=500, detail={"error": str(e)})
|
raise HTTPException(status_code=500, detail={"error": str(e)})
|
||||||
|
|
||||||
|
|
||||||
@router.post(
|
@router.post(
|
||||||
"/user/unblock",
|
"/end_user/unblock",
|
||||||
tags=["End User Management"],
|
tags=["End User Management"],
|
||||||
dependencies=[Depends(user_api_key_auth)],
|
dependencies=[Depends(user_api_key_auth)],
|
||||||
)
|
)
|
||||||
|
|
|
@ -6,6 +6,7 @@ import sys, os, asyncio, time, random
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
import traceback
|
import traceback
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
|
from fastapi import Request
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
import os
|
import os
|
||||||
|
@ -22,18 +23,87 @@ from litellm import Router, mock_completion
|
||||||
from litellm.proxy.utils import ProxyLogging
|
from litellm.proxy.utils import ProxyLogging
|
||||||
from litellm.proxy._types import UserAPIKeyAuth
|
from litellm.proxy._types import UserAPIKeyAuth
|
||||||
from litellm.caching import DualCache
|
from litellm.caching import DualCache
|
||||||
|
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
|
||||||
|
|
||||||
|
import pytest, logging, asyncio
|
||||||
|
import litellm, asyncio
|
||||||
|
from litellm.proxy.proxy_server import (
|
||||||
|
new_user,
|
||||||
|
generate_key_fn,
|
||||||
|
user_api_key_auth,
|
||||||
|
user_update,
|
||||||
|
delete_key_fn,
|
||||||
|
info_key_fn,
|
||||||
|
update_key_fn,
|
||||||
|
generate_key_fn,
|
||||||
|
generate_key_helper_fn,
|
||||||
|
spend_user_fn,
|
||||||
|
spend_key_fn,
|
||||||
|
view_spend_logs,
|
||||||
|
user_info,
|
||||||
|
block_user,
|
||||||
|
)
|
||||||
|
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
|
||||||
|
from litellm._logging import verbose_proxy_logger
|
||||||
|
|
||||||
|
verbose_proxy_logger.setLevel(level=logging.DEBUG)
|
||||||
|
|
||||||
|
from litellm.proxy._types import (
|
||||||
|
NewUserRequest,
|
||||||
|
GenerateKeyRequest,
|
||||||
|
DynamoDBArgs,
|
||||||
|
KeyRequest,
|
||||||
|
UpdateKeyRequest,
|
||||||
|
GenerateKeyRequest,
|
||||||
|
BlockUsers,
|
||||||
|
)
|
||||||
|
from litellm.proxy.utils import DBClient
|
||||||
|
from starlette.datastructures import URL
|
||||||
|
from litellm.caching import DualCache
|
||||||
|
|
||||||
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def prisma_client():
|
||||||
|
from litellm.proxy.proxy_cli import append_query_params
|
||||||
|
|
||||||
|
### add connection pool + pool timeout args
|
||||||
|
params = {"connection_limit": 100, "pool_timeout": 60}
|
||||||
|
database_url = os.getenv("DATABASE_URL")
|
||||||
|
modified_url = append_query_params(database_url, params)
|
||||||
|
os.environ["DATABASE_URL"] = modified_url
|
||||||
|
|
||||||
|
# Assuming DBClient is a class that needs to be instantiated
|
||||||
|
prisma_client = PrismaClient(
|
||||||
|
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
|
||||||
|
)
|
||||||
|
|
||||||
|
# Reset litellm.proxy.proxy_server.prisma_client to None
|
||||||
|
litellm.proxy.proxy_server.custom_db_client = None
|
||||||
|
litellm.proxy.proxy_server.litellm_proxy_budget_name = (
|
||||||
|
f"litellm-proxy-budget-{time.time()}"
|
||||||
|
)
|
||||||
|
litellm.proxy.proxy_server.user_custom_key_generate = None
|
||||||
|
|
||||||
|
return prisma_client
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_block_user_check():
|
async def test_block_user_check(prisma_client):
|
||||||
"""
|
"""
|
||||||
- Set a blocked user as a litellm module value
|
- Set a blocked user as a litellm module value
|
||||||
- Test to see if a call with that user id is made, an error is raised
|
- Test to see if a call with that user id is made, an error is raised
|
||||||
- Test to see if a call without that user is passes
|
- Test to see if a call without that user is passes
|
||||||
"""
|
"""
|
||||||
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
||||||
|
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||||
|
|
||||||
litellm.blocked_user_list = ["user_id_1"]
|
litellm.blocked_user_list = ["user_id_1"]
|
||||||
|
|
||||||
blocked_user_obj = _ENTERPRISE_BlockedUserList()
|
blocked_user_obj = _ENTERPRISE_BlockedUserList(
|
||||||
|
prisma_client=litellm.proxy.proxy_server.prisma_client
|
||||||
|
)
|
||||||
|
|
||||||
_api_key = "sk-12345"
|
_api_key = "sk-12345"
|
||||||
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
||||||
|
@ -61,3 +131,20 @@ async def test_block_user_check():
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
pytest.fail(f"An error occurred - {str(e)}")
|
pytest.fail(f"An error occurred - {str(e)}")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_block_user_db_check(prisma_client):
|
||||||
|
"""
|
||||||
|
- Block end user via "/user/block"
|
||||||
|
- Check returned value
|
||||||
|
"""
|
||||||
|
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
|
||||||
|
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
|
||||||
|
await litellm.proxy.proxy_server.prisma_client.connect()
|
||||||
|
_block_users = BlockUsers(user_ids=["user_id_1"])
|
||||||
|
result = await block_user(data=_block_users)
|
||||||
|
result = result["blocked_users"]
|
||||||
|
assert len(result) == 1
|
||||||
|
assert result[0].user_id == "user_id_1"
|
||||||
|
assert result[0].blocked == True
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue