Merge pull request #1588 from BerriAI/litellm_spend_per_user

[Feat] [WIP] View Spend Per User

docs/my-website/docs/proxy/ui.md

@@ -69,3 +69,7 @@ Connect your proxy to your UI, by entering:
 
 <Image img={require('../../img/spend_per_api_key.png')} />  
 
+### Spend Per User
+
+<Image img={require('../../img/spend_per_user.png')} />  
+

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(

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'"

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}")

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():