forked from phoenix/litellm-mirror
4d1b4beb3d
* use folder for caching * fix importing caching * fix clickhouse pyright * fix linting * fix correctly pass kwargs and args * fix test case for embedding * fix linting * fix embedding caching logic * fix refactor handle utils.py * fix test_embedding_caching_azure_individual_items_reordered
320 lines
10 KiB
Python
320 lines
10 KiB
Python
"""
|
|
Tests the following endpoints used by the UI
|
|
|
|
/global/spend/logs
|
|
/global/spend/keys
|
|
/global/spend/models
|
|
/global/activity
|
|
/global/activity/model
|
|
|
|
|
|
For all tests - test the following:
|
|
- Response is valid
|
|
- Response for Admin User is different from response from Internal User
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import traceback
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi import Request
|
|
from fastapi.routing import APIRoute
|
|
|
|
load_dotenv()
|
|
import io
|
|
import os
|
|
import time
|
|
|
|
# this file is to test litellm/proxy
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import asyncio
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm._logging import verbose_proxy_logger
|
|
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
|
new_user,
|
|
user_info,
|
|
user_update,
|
|
)
|
|
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
|
delete_key_fn,
|
|
generate_key_fn,
|
|
generate_key_helper_fn,
|
|
info_key_fn,
|
|
regenerate_key_fn,
|
|
update_key_fn,
|
|
)
|
|
from litellm.proxy.management_endpoints.team_endpoints import (
|
|
new_team,
|
|
team_info,
|
|
update_team,
|
|
)
|
|
from litellm.proxy.proxy_server import (
|
|
LitellmUserRoles,
|
|
audio_transcriptions,
|
|
chat_completion,
|
|
completion,
|
|
embeddings,
|
|
image_generation,
|
|
model_list,
|
|
moderations,
|
|
new_end_user,
|
|
user_api_key_auth,
|
|
)
|
|
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
|
global_spend,
|
|
global_spend_logs,
|
|
global_spend_models,
|
|
global_spend_keys,
|
|
spend_key_fn,
|
|
spend_user_fn,
|
|
view_spend_logs,
|
|
)
|
|
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend
|
|
|
|
verbose_proxy_logger.setLevel(level=logging.DEBUG)
|
|
|
|
from starlette.datastructures import URL
|
|
|
|
from litellm.caching.caching import DualCache
|
|
from litellm.proxy._types import (
|
|
DynamoDBArgs,
|
|
GenerateKeyRequest,
|
|
RegenerateKeyRequest,
|
|
KeyRequest,
|
|
LiteLLM_UpperboundKeyGenerateParams,
|
|
NewCustomerRequest,
|
|
NewTeamRequest,
|
|
NewUserRequest,
|
|
ProxyErrorTypes,
|
|
ProxyException,
|
|
UpdateKeyRequest,
|
|
UpdateTeamRequest,
|
|
UpdateUserRequest,
|
|
UserAPIKeyAuth,
|
|
)
|
|
|
|
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 PrismaClient 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.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()
|
|
async def test_view_daily_spend_ui(prisma_client):
|
|
print("prisma client=", prisma_client)
|
|
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()
|
|
from litellm.proxy.proxy_server import user_api_key_cache
|
|
|
|
spend_logs_for_admin = await global_spend_logs(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
),
|
|
api_key=None,
|
|
)
|
|
|
|
print("spend_logs_for_admin=", spend_logs_for_admin)
|
|
|
|
spend_logs_for_internal_user = await global_spend_logs(
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
|
|
),
|
|
api_key=None,
|
|
)
|
|
|
|
print("spend_logs_for_internal_user=", spend_logs_for_internal_user)
|
|
|
|
# Calculate total spend for admin
|
|
admin_total_spend = sum(log.get("spend", 0) for log in spend_logs_for_admin)
|
|
|
|
# Calculate total spend for internal user (0 in this case, but we'll keep it generic)
|
|
internal_user_total_spend = sum(
|
|
log.get("spend", 0) for log in spend_logs_for_internal_user
|
|
)
|
|
|
|
print("total_spend_for_admin=", admin_total_spend)
|
|
print("total_spend_for_internal_user=", internal_user_total_spend)
|
|
|
|
assert (
|
|
admin_total_spend > internal_user_total_spend
|
|
), "Admin should have more spend than internal user"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_global_spend_models(prisma_client):
|
|
print("prisma client=", prisma_client)
|
|
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()
|
|
|
|
# Test for admin user
|
|
models_spend_for_admin = await global_spend_models(
|
|
limit=10,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
),
|
|
)
|
|
|
|
print("models_spend_for_admin=", models_spend_for_admin)
|
|
|
|
# Test for internal user
|
|
models_spend_for_internal_user = await global_spend_models(
|
|
limit=10,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
|
|
),
|
|
)
|
|
|
|
print("models_spend_for_internal_user=", models_spend_for_internal_user)
|
|
|
|
# Assertions
|
|
assert isinstance(models_spend_for_admin, list), "Admin response should be a list"
|
|
assert isinstance(
|
|
models_spend_for_internal_user, list
|
|
), "Internal user response should be a list"
|
|
|
|
# Check if the response has the expected shape for both admin and internal user
|
|
expected_keys = ["model", "total_spend"]
|
|
|
|
if len(models_spend_for_admin) > 0:
|
|
assert all(
|
|
key in models_spend_for_admin[0] for key in expected_keys
|
|
), f"Admin response should contain keys: {expected_keys}"
|
|
assert isinstance(
|
|
models_spend_for_admin[0]["model"], str
|
|
), "Model should be a string"
|
|
assert isinstance(
|
|
models_spend_for_admin[0]["total_spend"], (int, float)
|
|
), "Total spend should be a number"
|
|
|
|
if len(models_spend_for_internal_user) > 0:
|
|
assert all(
|
|
key in models_spend_for_internal_user[0] for key in expected_keys
|
|
), f"Internal user response should contain keys: {expected_keys}"
|
|
assert isinstance(
|
|
models_spend_for_internal_user[0]["model"], str
|
|
), "Model should be a string"
|
|
assert isinstance(
|
|
models_spend_for_internal_user[0]["total_spend"], (int, float)
|
|
), "Total spend should be a number"
|
|
|
|
# Check if the lists are sorted by total_spend in descending order
|
|
if len(models_spend_for_admin) > 1:
|
|
assert all(
|
|
models_spend_for_admin[i]["total_spend"]
|
|
>= models_spend_for_admin[i + 1]["total_spend"]
|
|
for i in range(len(models_spend_for_admin) - 1)
|
|
), "Admin response should be sorted by total_spend in descending order"
|
|
|
|
if len(models_spend_for_internal_user) > 1:
|
|
assert all(
|
|
models_spend_for_internal_user[i]["total_spend"]
|
|
>= models_spend_for_internal_user[i + 1]["total_spend"]
|
|
for i in range(len(models_spend_for_internal_user) - 1)
|
|
), "Internal user response should be sorted by total_spend in descending order"
|
|
|
|
# Check if admin has access to more or equal models compared to internal user
|
|
assert len(models_spend_for_admin) >= len(
|
|
models_spend_for_internal_user
|
|
), "Admin should have access to at least as many models as internal user"
|
|
|
|
# Check if the response contains expected fields
|
|
if len(models_spend_for_admin) > 0:
|
|
assert all(
|
|
key in models_spend_for_admin[0] for key in ["model", "total_spend"]
|
|
), "Admin response should contain model, total_spend, and total_tokens"
|
|
|
|
if len(models_spend_for_internal_user) > 0:
|
|
assert all(
|
|
key in models_spend_for_internal_user[0] for key in ["model", "total_spend"]
|
|
), "Internal user response should contain model, total_spend, and total_tokens"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_global_spend_keys(prisma_client):
|
|
print("prisma client=", prisma_client)
|
|
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()
|
|
|
|
# Test for admin user
|
|
keys_spend_for_admin = await global_spend_keys(
|
|
limit=10,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
),
|
|
)
|
|
|
|
print("keys_spend_for_admin=", keys_spend_for_admin)
|
|
|
|
# Test for internal user
|
|
keys_spend_for_internal_user = await global_spend_keys(
|
|
limit=10,
|
|
user_api_key_dict=UserAPIKeyAuth(
|
|
api_key="sk-1234", user_role=LitellmUserRoles.INTERNAL_USER, user_id="1234"
|
|
),
|
|
)
|
|
|
|
print("keys_spend_for_internal_user=", keys_spend_for_internal_user)
|
|
|
|
# Assertions
|
|
assert isinstance(keys_spend_for_admin, list), "Admin response should be a list"
|
|
assert isinstance(
|
|
keys_spend_for_internal_user, list
|
|
), "Internal user response should be a list"
|
|
|
|
# Check if admin has access to more or equal keys compared to internal user
|
|
assert len(keys_spend_for_admin) >= len(
|
|
keys_spend_for_internal_user
|
|
), "Admin should have access to at least as many keys as internal user"
|
|
|
|
# Check if the response contains expected fields
|
|
if len(keys_spend_for_admin) > 0:
|
|
assert all(
|
|
key in keys_spend_for_admin[0]
|
|
for key in ["api_key", "total_spend", "key_alias", "key_name"]
|
|
), "Admin response should contain api_key, total_spend, key_alias, and key_name"
|
|
|
|
if len(keys_spend_for_internal_user) > 0:
|
|
assert all(
|
|
key in keys_spend_for_internal_user[0]
|
|
for key in ["api_key", "total_spend", "key_alias", "key_name"]
|
|
), "Internal user response should contain api_key, total_spend, key_alias, and key_name"
|