diff --git a/docs/my-website/docs/proxy/ui.md b/docs/my-website/docs/proxy/ui.md index 0a3f59017..7b4c6ed47 100644 --- a/docs/my-website/docs/proxy/ui.md +++ b/docs/my-website/docs/proxy/ui.md @@ -69,3 +69,7 @@ Connect your proxy to your UI, by entering: +### Spend Per User + + + diff --git a/docs/my-website/img/spend_per_user.png b/docs/my-website/img/spend_per_user.png new file mode 100644 index 000000000..066c4baaf Binary files /dev/null and b/docs/my-website/img/spend_per_user.png differ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cc4222f8e..ca58371f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2393,7 +2393,7 @@ async def info_key_fn( @router.get( "/spend/keys", - tags=["Budget & Spend Tracking"], + tags=["budget & spend Tracking"], dependencies=[Depends(user_api_key_auth)], ) async def spend_key_fn(): @@ -2424,9 +2424,44 @@ async def spend_key_fn(): ) +@router.get( + "/spend/users", + tags=["budget & spend Tracking"], + dependencies=[Depends(user_api_key_auth)], +) +async def spend_user_fn(): + """ + View all users created, ordered by spend + + Example Request: + ``` + curl -X GET "http://0.0.0.0:8000/spend/users" \ +-H "Authorization: Bearer sk-1234" + ``` + """ + global prisma_client + try: + if prisma_client is None: + raise Exception( + f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + ) + + user_info = await prisma_client.get_data( + table_name="user", query_type="find_all" + ) + + return user_info + + except Exception as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": str(e)}, + ) + + @router.get( "/spend/logs", - tags=["Budget & Spend Tracking"], + tags=["budget & spend Tracking"], dependencies=[Depends(user_api_key_auth)], ) async def view_spend_logs( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1ac6f7f06..978355568 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -432,6 +432,11 @@ class PrismaClient: } ) return response + elif table_name == "user" and query_type == "find_all": + response = await self.db.litellm_usertable.find_many( # type: ignore + order={"spend": "desc"}, + ) + return response elif table_name == "spend": verbose_proxy_logger.debug( f"PrismaClient: get_data: table_name == 'spend'" diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index f130807c6..49f091cd6 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -42,6 +42,8 @@ from litellm.proxy.proxy_server import ( info_key_fn, update_key_fn, generate_key_fn, + spend_user_fn, + spend_key_fn, view_spend_logs, ) from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -851,3 +853,39 @@ async def test_call_with_key_over_budget_stream(prisma_client): error_detail = e.message assert "Authentication Error, ExceededTokenBudget:" in error_detail print(vars(e)) + + +@pytest.mark.asyncio() +async def test_view_spend_per_user(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() + try: + user_by_spend = await spend_user_fn() + assert type(user_by_spend) == list + assert len(user_by_spend) > 0 + first_user = user_by_spend[0] + + print("\nfirst_user=", first_user) + assert first_user.spend > 0 + except Exception as e: + print("Got Exception", e) + pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio() +async def test_view_spend_per_key(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() + try: + key_by_spend = await spend_key_fn() + assert type(key_by_spend) == list + assert len(key_by_spend) > 0 + first_key = key_by_spend[0] + + print("\nfirst_key=", first_key) + assert first_key.spend > 0 + except Exception as e: + print("Got Exception", e) + pytest.fail(f"Got exception {e}") diff --git a/ui/admin.py b/ui/admin.py index 8b5c6b3ab..96da791df 100644 --- a/ui/admin.py +++ b/ui/admin.py @@ -238,7 +238,59 @@ def spend_per_key(): def spend_per_user(): - pass + import streamlit as st + import requests + + # Check if the necessary configuration is available + if ( + st.session_state.get("api_url", None) is not None + and st.session_state.get("proxy_key", None) is not None + ): + # Make the GET request + try: + complete_url = "" + if isinstance(st.session_state["api_url"], str) and st.session_state[ + "api_url" + ].endswith("/"): + complete_url = f"{st.session_state['api_url']}/spend/users" + else: + complete_url = f"{st.session_state['api_url']}/spend/users" + response = requests.get( + complete_url, + headers={"Authorization": f"Bearer {st.session_state['proxy_key']}"}, + ) + # Check if the request was successful + if response.status_code == 200: + spend_per_user = response.json() + # Create DataFrame + spend_df = pd.DataFrame(spend_per_user) + + # Display the spend per key as a graph + st.header("Spend ($) per User:") + top_10_df = spend_df.nlargest(10, "spend") + fig = px.bar( + top_10_df, + x="user_id", + y="spend", + title="Top 10 Spend per User", + height=550, # Adjust the height + width=1200, # Adjust the width) + hover_data=["user_id", "spend", "max_budget"], + ) + st.plotly_chart(fig) + + # Display the spend per key as a table + st.write("Spend per User - Full Table:") + st.table(spend_df) + + else: + st.error(f"Failed to get models. Status code: {response.status_code}") + except Exception as e: + st.error(f"An error occurred while requesting models: {e}") + else: + st.warning( + f"Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page. Currently set Proxy Endpoint: {st.session_state.get('api_url', None)} and Proxy Key: {st.session_state.get('proxy_key', None)}" + ) def create_key():