diff --git a/.circleci/config.yml b/.circleci/config.yml index 5cd3c50d504..dccaa2b1122 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -111,6 +111,44 @@ jobs: # Store test results - store_test_results: path: test-results + ui_endpoint_testing: + docker: + - image: cimg/python:3.11 + working_directory: ~/project + + steps: + - checkout + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-asyncio==0.21.1" + - save_cache: + paths: + - ./venv + key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }} + - run: + name: Run prisma ./entrypoint.sh + command: | + set +e + chmod +x entrypoint.sh + ./entrypoint.sh + set -e + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/proxy_admin_ui_tests -x --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + + # Store test results + - store_test_results: + path: test-results installing_litellm_on_python: docker: @@ -251,7 +289,7 @@ jobs: command: | pwd ls - python -m pytest -s -vv tests/ -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests + python -m pytest -s -vv tests/ -x --junitxml=test-results/junit.xml --durations=5 --ignore=tests/otel_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests no_output_timeout: 120m # Store test results @@ -539,6 +577,12 @@ workflows: only: - main - /litellm_.*/ + - ui_endpoint_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - build_and_test: filters: branches: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00f0cb7e300..67acf71e52d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -344,6 +344,13 @@ class LiteLLMRoutes(enum.Enum): "/key/update", "/key/delete", "/key/info", + "/global/spend/tags", + "/global/spend/keys", + "/global/spend/models", + "/global/spend/provider", + "/global/spend/end_users", + "/global/activity", + "/global/activity/model", ] + spend_tracking_routes + sso_only_routes diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 28db48d7ff8..86f95936279 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -177,6 +177,35 @@ async def view_spend_tags( ) +async def get_global_activity_internal_user( + user_api_key_dict: UserAPIKeyAuth, start_date: datetime, end_date: datetime +): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + user_id = user_api_key_dict.user_id + if user_id is None: + raise HTTPException(status_code=500, detail={"error": "No user_id found"}) + + sql_query = """ + SELECT + date_trunc('day', "startTime") AS date, + COUNT(*) AS api_requests, + SUM(total_tokens) AS total_tokens + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + AND "user" = $3 + GROUP BY date_trunc('day', "startTime") + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date, end_date, user_id + ) + + return db_response + + @router.get( "/global/activity", tags=["Budget & Spend Tracking"], @@ -195,6 +224,7 @@ async def get_global_activity( default=None, description="Time till which to view spend", ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Get number of API Requests, total tokens through proxy @@ -236,18 +266,27 @@ async def get_global_activity( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - sql_query = """ - SELECT - date_trunc('day', "startTime") AS date, - COUNT(*) AS api_requests, - SUM(total_tokens) AS total_tokens - FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' - GROUP BY date_trunc('day', "startTime") - """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + db_response = await get_global_activity_internal_user( + user_api_key_dict, start_date_obj, end_date_obj + ) + else: + + sql_query = """ + SELECT + date_trunc('day', "startTime") AS date, + COUNT(*) AS api_requests, + SUM(total_tokens) AS total_tokens + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + GROUP BY date_trunc('day', "startTime") + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj + ) if db_response is None: return [] @@ -282,6 +321,36 @@ async def get_global_activity( ) +async def get_global_activity_model_internal_user( + user_api_key_dict: UserAPIKeyAuth, start_date: datetime, end_date: datetime +): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + user_id = user_api_key_dict.user_id + if user_id is None: + raise HTTPException(status_code=500, detail={"error": "No user_id found"}) + + sql_query = """ + SELECT + model_group, + date_trunc('day', "startTime") AS date, + COUNT(*) AS api_requests, + SUM(total_tokens) AS total_tokens + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + AND "user" = $3 + GROUP BY model_group, date_trunc('day', "startTime") + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date, end_date, user_id + ) + + return db_response + + @router.get( "/global/activity/model", tags=["Budget & Spend Tracking"], @@ -300,6 +369,7 @@ async def get_global_activity_model( default=None, description="Time till which to view spend", ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Get number of API Requests, total tokens through proxy - Grouped by MODEL @@ -364,19 +434,28 @@ async def get_global_activity_model( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - sql_query = """ - SELECT - model_group, - date_trunc('day', "startTime") AS date, - COUNT(*) AS api_requests, - SUM(total_tokens) AS total_tokens - FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' - GROUP BY model_group, date_trunc('day', "startTime") - """ - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + db_response = await get_global_activity_model_internal_user( + user_api_key_dict, start_date_obj, end_date_obj + ) + else: + + sql_query = """ + SELECT + model_group, + date_trunc('day', "startTime") AS date, + COUNT(*) AS api_requests, + SUM(total_tokens) AS total_tokens + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date + interval '1 day' + GROUP BY model_group, date_trunc('day', "startTime") + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj + ) if db_response is None: return [] @@ -711,6 +790,7 @@ async def get_global_spend_provider( default=None, description="Time till which to view spend", ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Get breakdown of spend per provider @@ -748,19 +828,42 @@ async def get_global_spend_provider( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - sql_query = """ + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + user_id = user_api_key_dict.user_id + if user_id is None: + raise HTTPException( + status_code=400, detail={"error": "No user_id found"} + ) - SELECT - model_id, - SUM(spend) AS spend - FROM "LiteLLM_SpendLogs" - WHERE "startTime" BETWEEN $1::date AND $2::date AND length(model_id) > 0 - GROUP BY model_id - """ + sql_query = """ + SELECT + model_id, + SUM(spend) AS spend + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date + AND length(model_id) > 0 + AND "user" = $3 + GROUP BY model_id + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj, user_id + ) + else: + sql_query = """ + SELECT + model_id, + SUM(spend) AS spend + FROM "LiteLLM_SpendLogs" + WHERE "startTime" BETWEEN $1::date AND $2::date AND length(model_id) > 0 + GROUP BY model_id + """ + db_response = await prisma_client.db.query_raw( + sql_query, start_date_obj, end_date_obj + ) - db_response = await prisma_client.db.query_raw( - sql_query, start_date_obj, end_date_obj - ) if db_response is None: return [] @@ -1515,6 +1618,7 @@ async def view_spend_logs( default=None, description="Time till which to view key spend", ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ View all spend logs, if request_id is provided, only logs for that request_id will be returned @@ -1545,6 +1649,12 @@ async def view_spend_logs( """ from litellm.proxy.proxy_server import prisma_client + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + user_id = user_api_key_dict.user_id + try: verbose_proxy_logger.debug("inside view_spend_logs") if prisma_client is None: @@ -1733,6 +1843,45 @@ async def global_spend_reset(): } +async def global_spend_for_internal_user( + api_key: Optional[str] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise ProxyException( + message="Prisma Client is not initialized", + type="internal_error", + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + try: + + user_id = user_api_key_dict.user_id + if user_id is None: + raise ValueError(f"/global/spend/logs Error: User ID is None") + if api_key is not None: + sql_query = """ + SELECT * FROM "MonthlyGlobalSpendPerUserPerKey" + WHERE "api_key" = $1 AND "user" = $2 + ORDER BY "date"; + """ + + response = await prisma_client.db.query_raw(sql_query, api_key, user_id) + + return response + + sql_query = """SELECT * FROM "MonthlyGlobalSpendPerUserPerKey" WHERE "user" = $1 ORDER BY "date";""" + + response = await prisma_client.db.query_raw(sql_query, user_id) + + return response + except Exception as e: + verbose_proxy_logger.error(f"/global/spend/logs Error: {str(e)}") + raise e + + @router.get( "/global/spend/logs", tags=["Budget & Spend Tracking"], @@ -1743,7 +1892,8 @@ async def global_spend_logs( api_key: str = fastapi.Query( default=None, description="API Key to get global spend (spend per day for last 30d). Admin-only endpoint", - ) + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ [BETA] This is a beta endpoint. It will change. @@ -1764,6 +1914,17 @@ async def global_spend_logs( param="None", code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + response = await global_spend_for_internal_user( + api_key=api_key, user_api_key_dict=user_api_key_dict + ) + + return response + if api_key is None: sql_query = """SELECT * FROM "MonthlyGlobalSpend" ORDER BY "date";""" @@ -1784,6 +1945,7 @@ async def global_spend_logs( except Exception as e: error_trace = traceback.format_exc() error_str = str(e) + "\n" + error_trace + verbose_proxy_logger.error(f"/global/spend/logs Error: {error_str}") if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", f"/global/spend/logs Error({error_str})"), @@ -1851,6 +2013,52 @@ async def global_spend(): ) +async def global_spend_key_internal_user( + user_api_key_dict: UserAPIKeyAuth, limit: int = 10 +): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + user_id = user_api_key_dict.user_id + if user_id is None: + raise HTTPException(status_code=500, detail={"error": "No user_id found"}) + + sql_query = """ + WITH top_api_keys AS ( + SELECT + api_key, + SUM(spend) as total_spend + FROM + "LiteLLM_SpendLogs" + WHERE + "user" = $1 + GROUP BY + api_key + ORDER BY + total_spend DESC + LIMIT $2 -- Adjust this number to get more or fewer top keys + ) + SELECT + t.api_key, + t.total_spend, + v.key_alias, + v.key_name + FROM + top_api_keys t + LEFT JOIN + "LiteLLM_VerificationToken" v ON t.api_key = v.token + ORDER BY + t.total_spend DESC; + + """ + + response = await prisma_client.db.query_raw(sql_query, user_id, limit) + + return response + + @router.get( "/global/spend/keys", tags=["Budget & Spend Tracking"], @@ -1861,7 +2069,8 @@ async def global_spend_keys( limit: int = fastapi.Query( default=None, description="Number of keys to get. Will return Top 'n' keys.", - ) + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ [BETA] This is a beta endpoint. It will change. @@ -1870,6 +2079,15 @@ async def global_spend_keys( """ from litellm.proxy.proxy_server import prisma_client + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + response = await global_spend_key_internal_user( + user_api_key_dict=user_api_key_dict + ) + + return response if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) sql_query = f"""SELECT * FROM "Last30dKeysBySpend" LIMIT {limit};""" @@ -2057,6 +2275,39 @@ LIMIT 100 return response +async def global_spend_models_internal_user( + user_api_key_dict: UserAPIKeyAuth, limit: int = 10 +): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + user_id = user_api_key_dict.user_id + if user_id is None: + raise HTTPException(status_code=500, detail={"error": "No user_id found"}) + + sql_query = """ + SELECT + model, + SUM(spend) as total_spend, + SUM(total_tokens) as total_tokens + FROM + "LiteLLM_SpendLogs" + WHERE + "user" = $1 + GROUP BY + model + ORDER BY + total_spend DESC + LIMIT $2; + """ + + response = await prisma_client.db.query_raw(sql_query, user_id, limit) + + return response + + @router.get( "/global/spend/models", tags=["Budget & Spend Tracking"], @@ -2065,17 +2316,27 @@ LIMIT 100 ) async def global_spend_models( limit: int = fastapi.Query( - default=None, + default=10, description="Number of models to get. Will return Top 'n' models.", - ) + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ [BETA] This is a beta endpoint. It will change. - Use this to get the top 'n' keys with the highest spend, ordered by spend. + Use this to get the top 'n' models with the highest spend, ordered by spend. """ from litellm.proxy.proxy_server import prisma_client + if ( + user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER + or user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ): + response = await global_spend_models_internal_user( + user_api_key_dict=user_api_key_dict, limit=limit + ) + return response + if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2423fb105a9..0d6edd7702e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -969,11 +969,12 @@ class PrismaClient: 'Last30dKeysBySpend', 'Last30dModelsBySpend', 'MonthlyGlobalSpendPerKey', + 'MonthlyGlobalSpendPerUserPerKey', 'Last30dTopEndUsersSpend' ) """ ) - if ret[0]["sum"] == 6: + if ret[0]["sum"] == 7: print("All necessary views exist!") # noqa return except Exception: @@ -1097,6 +1098,31 @@ class PrismaClient: await self.db.execute_raw(query=sql_query) print("MonthlyGlobalSpendPerKey Created!") # noqa + try: + await self.db.query_raw( + """SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""" + ) + print("MonthlyGlobalSpendPerUserPerKey Exists!") # noqa + except Exception as e: + sql_query = """ + CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerUserPerKey" AS + SELECT + DATE("startTime") AS date, + SUM("spend") AS spend, + api_key as api_key, + "user" as "user" + FROM + "LiteLLM_SpendLogs" + WHERE + "startTime" >= (CURRENT_DATE - INTERVAL '20 days') + GROUP BY + DATE("startTime"), + "user", + api_key; + """ + await self.db.execute_raw(query=sql_query) + + print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa try: await self.db.query_raw( diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index a4a6606c48e..73f6264c603 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -11,7 +11,8 @@ import os sys.path.insert( 0, os.path.abspath("../..") -) # Adds the parent directory to the system-path +) # Adds the parent directory to the systempath + import os from unittest.mock import AsyncMock, MagicMock, patch @@ -23,7 +24,7 @@ from litellm import RateLimitError, Timeout, completion, completion_cost, embedd from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.prompt_templates.factory import anthropic_messages_pt -# litellm.num_retries=3 +# litellm.num_retries = 3 litellm.cache = None litellm.success_callback = [] user_message = "Write a short poem about the sky" diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 708025d1d6d..adf0e8aea96 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -537,6 +537,7 @@ def test_call_with_user_over_budget(prisma_client): asyncio.run(test()) except Exception as e: + print("got an errror=", e) error_detail = e.message assert "ExceededBudget:" in error_detail assert isinstance(e, ProxyException) @@ -1492,7 +1493,10 @@ def test_call_with_key_over_budget(prisma_client): proxy_logging_obj=proxy_logging_obj, ) # test spend_log was written and we can read it - spend_logs = await view_spend_logs(request_id=request_id) + spend_logs = await view_spend_logs( + request_id=request_id, + user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), + ) print("read spend logs", spend_logs) assert len(spend_logs) == 1 @@ -1607,7 +1611,10 @@ def test_call_with_key_over_budget_no_cache(prisma_client): proxy_logging_obj=proxy_logging_obj, ) # test spend_log was written and we can read it - spend_logs = await view_spend_logs(request_id=request_id) + spend_logs = await view_spend_logs( + request_id=request_id, + user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), + ) print("read spend logs", spend_logs) assert len(spend_logs) == 1 @@ -1727,7 +1734,10 @@ def test_call_with_key_over_model_budget(prisma_client): proxy_logging_obj=proxy_logging_obj, ) # test spend_log was written and we can read it - spend_logs = await view_spend_logs(request_id=request_id) + spend_logs = await view_spend_logs( + request_id=request_id, + user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), + ) print("read spend logs", spend_logs) assert len(spend_logs) == 1 @@ -2296,7 +2306,10 @@ async def test_proxy_load_test_db(prisma_client): await asyncio.sleep(120) try: # call spend logs - spend_logs = await view_spend_logs(api_key=generated_key) + spend_logs = await view_spend_logs( + api_key=generated_key, + user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), + ) print(f"len responses: {len(spend_logs)}") assert len(spend_logs) == n diff --git a/tests/proxy_admin_ui_tests/test_usage_endpoints.py b/tests/proxy_admin_ui_tests/test_usage_endpoints.py new file mode 100644 index 00000000000..9918015a1c6 --- /dev/null +++ b/tests/proxy_admin_ui_tests/test_usage_endpoints.py @@ -0,0 +1,321 @@ +""" +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 import DualCache +from litellm.proxy._types import ( + DynamoDBArgs, + GenerateKeyRequest, + KeyRequest, + LiteLLM_UpperboundKeyGenerateParams, + NewCustomerRequest, + NewTeamRequest, + NewUserRequest, + ProxyErrorTypes, + ProxyException, + UpdateKeyRequest, + UpdateTeamRequest, + UpdateUserRequest, + UserAPIKeyAuth, +) +from litellm.proxy.utils import DBClient + +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() +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" diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index c8f5745ed49..728a35076dd 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -12,6 +12,8 @@ interface SidebarProps { defaultSelectedKey: string[] | null; } +const rolesAllowedToSeeUsage = ["Admin", "Admin Viewer", "Internal User", "Internal Viewer"]; + const Sidebar: React.FC = ({ setPage, userRole, @@ -62,7 +64,7 @@ const Sidebar: React.FC = ({ Models ) : null} - {userRole == "Admin" ? ( + {rolesAllowedToSeeUsage.includes(userRole) ? ( setPage("usage")}> Usage diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 23c64d4373b..fbe43aa12f2 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/ui/litellm-dashboard/src/components/usage.tsx @@ -132,6 +132,13 @@ type DataDict = { [key: string]: unknown }; type UserData = { user_id: string; spend: number }; +const isAdminOrAdminViewer = (role: string | null): boolean => { + if (role === null) return false; + return role === 'Admin' || role === 'Admin Viewer'; +}; + + + const UsagePage: React.FC = ({ accessToken, token, @@ -256,125 +263,144 @@ const UsagePage: React.FC = ({ const valueFormatter = (number: number) => `$ ${new Intl.NumberFormat("us").format(number).toString()}`; + const fetchAndSetData = async ( + fetchFunction: () => Promise, + setStateFunction: React.Dispatch>, + errorMessage: string + ) => { + try { + const data = await fetchFunction(); + setStateFunction(data); + } catch (error) { + console.error(errorMessage, error); + // Optionally, update UI to reflect error state for this specific data + } + }; + + const fetchOverallSpend = () => fetchAndSetData( + () => accessToken ? adminSpendLogsCall(accessToken) : Promise.reject("No access token"), + setKeySpendData, + "Error fetching overall spend" + ); + + const fetchProviderSpend = () => fetchAndSetData( + () => accessToken && token ? adminspendByProvider(accessToken, token, startTime, endTime) : Promise.reject("No access token or token"), + setSpendByProvider, + "Error fetching provider spend" + ); + + const fetchTopKeys = async () => { + if (!accessToken) return; + await fetchAndSetData( + async () => { + const top_keys = await adminTopKeysCall(accessToken); + return top_keys.map((k: any) => ({ + key: (k["key_alias"] || k["key_name"] || k["api_key"]).substring(0, 10), + spend: k["total_spend"], + })); + }, + setTopKeys, + "Error fetching top keys" + ); + }; + + const fetchTopModels = async () => { + if (!accessToken) return; + await fetchAndSetData( + async () => { + const top_models = await adminTopModelsCall(accessToken); + return top_models.map((k: any) => ({ + key: k["model"], + spend: k["total_spend"], + })); + }, + setTopModels, + "Error fetching top models" + ); + }; + + const fetchTeamSpend = async () => { + if (!accessToken) return; + await fetchAndSetData( + async () => { + const teamSpend = await teamSpendLogsCall(accessToken); + setTeamSpendData(teamSpend.daily_spend); + setUniqueTeamIds(teamSpend.teams); + return teamSpend.total_spend_per_team.map((tspt: any) => ({ + name: tspt["team_id"] || "", + value: (tspt["total_spend"] || 0).toFixed(2), + })); + }, + setTotalSpendPerTeam, + "Error fetching team spend" + ); + }; + + const fetchTagNames = () => { + if (!accessToken) return; + fetchAndSetData( + async () => { + const all_tag_names = await allTagNamesCall(accessToken); + return all_tag_names.tag_names; + }, + setAllTagNames, + "Error fetching tag names" + ); + }; + + const fetchTopTags = () => { + if (!accessToken) return; + fetchAndSetData( + () => tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined), + (data) => setTopTagsData(data.spend_per_tag), + "Error fetching top tags" + ); + }; + + const fetchTopEndUsers = () => { + if (!accessToken) return; + fetchAndSetData( + () => adminTopEndUsersCall(accessToken, null, undefined, undefined), + setTopUsers, + "Error fetching top end users" + ); + }; + + const fetchGlobalActivity = () => { + if (!accessToken) return; + fetchAndSetData( + () => adminGlobalActivity(accessToken, startTime, endTime), + setGlobalActivity, + "Error fetching global activity" + ); + }; + + const fetchGlobalActivityPerModel = () => { + if (!accessToken) return; + fetchAndSetData( + () => adminGlobalActivityPerModel(accessToken, startTime, endTime), + setGlobalActivityPerModel, + "Error fetching global activity per model" + ); + }; + useEffect(() => { if (accessToken && token && userRole && userID) { - const fetchData = async () => { - try { - /** - * If user is Admin - query the global views endpoints - * If user is App Owner - use the normal spend logs call - */ - console.log(`user role: ${userRole}`); - if (userRole == "Admin" || userRole == "Admin Viewer") { - const overall_spend = await adminSpendLogsCall(accessToken); - setKeySpendData(overall_spend); - - const provider_spend = await adminspendByProvider(accessToken, token, startTime, endTime); - console.log("provider_spend", provider_spend); - setSpendByProvider(provider_spend); - const top_keys = await adminTopKeysCall(accessToken); - const filtered_keys = top_keys.map((k: any) => ({ - key: (k["key_alias"] || k["key_name"] || k["api_key"]).substring( - 0, - 10 - ), - spend: k["total_spend"], - })); - setTopKeys(filtered_keys); - const top_models = await adminTopModelsCall(accessToken); - const filtered_models = top_models.map((k: any) => ({ - key: k["model"], - spend: k["total_spend"], - })); - setTopModels(filtered_models); + fetchOverallSpend(); + fetchProviderSpend(); + fetchTopKeys(); + fetchTopModels(); + fetchGlobalActivity(); + fetchGlobalActivityPerModel(); - const teamSpend = await teamSpendLogsCall(accessToken); - console.log("teamSpend", teamSpend); - setTeamSpendData(teamSpend.daily_spend); - setUniqueTeamIds(teamSpend.teams) - - let total_spend_per_team = teamSpend.total_spend_per_team; - // in total_spend_per_team, replace null team_id with "" and replace null total_spend with 0 - - total_spend_per_team = total_spend_per_team.map((tspt: any) => { - tspt["name"] = tspt["team_id"] || ""; - tspt["value"] = tspt["total_spend"] || 0; - // round the value to 2 decimal places - - tspt["value"] = tspt["value"].toFixed(2); - - - return tspt; - }) - - setTotalSpendPerTeam(total_spend_per_team); - - // all_tag_names -> used for dropdown - const all_tag_names = await allTagNamesCall(accessToken); - setAllTagNames(all_tag_names.tag_names); - - //get top tags - const top_tags = await tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined); - setTopTagsData(top_tags.spend_per_tag); - - - // get spend per end-user - let spend_user_call = await adminTopEndUsersCall(accessToken, null, undefined, undefined); - setTopUsers(spend_user_call); - - console.log("spend/user result", spend_user_call); - - let global_activity_response = await adminGlobalActivity(accessToken, startTime, endTime); - setGlobalActivity(global_activity_response) - - let global_activity_per_model = await adminGlobalActivityPerModel(accessToken, startTime, endTime); - console.log("global activity per model", global_activity_per_model); - setGlobalActivityPerModel(global_activity_per_model) - - - } else if (userRole == "App Owner") { - await userSpendLogsCall( - accessToken, - token, - userRole, - userID, - startTime, - endTime - ).then(async (response) => { - console.log("result from spend logs call", response); - if ("daily_spend" in response) { - // this is from clickhouse analytics - // - let daily_spend = response["daily_spend"]; - console.log("daily spend", daily_spend); - setKeySpendData(daily_spend); - let topApiKeys = response.top_api_keys; - setTopKeys(topApiKeys); - } else { - const topKeysResponse = await keyInfoCall( - accessToken, - getTopKeys(response) - ); - const filtered_keys = topKeysResponse["info"].map((k: any) => ({ - key: ( - k["key_name"] || - k["key_alias"] - ).substring(0, 10), - spend: k["spend"], - })); - setTopKeys(filtered_keys); - setKeySpendData(response); - } - }); - } - } catch (error) { - console.error("There was an error fetching the data", error); - // Optionally, update your UI to reflect the error state here as well - } - }; - fetchData(); + if (isAdminOrAdminViewer(userRole)) { + fetchTeamSpend(); + fetchTagNames(); + fetchTopTags(); + fetchTopEndUsers(); + } } }, [accessToken, token, userRole, userID, startTime, endTime]); @@ -385,9 +411,17 @@ const UsagePage: React.FC = ({ All Up - Team Based Usage - Customer Usage - Tag Based Usage + + {isAdminOrAdminViewer(userRole) ? ( + <> + Team Based Usage + Customer Usage + Tag Based Usage + + ) : ( + <>
+ + )}