Merge pull request #1575 from BerriAI/litellm_ui_view_spend

[Feat] UI view spend per API Key

litellm/proxy/admin_ui.py

@@ -98,7 +98,7 @@ def list_models():
             st.error(f"An error occurred while requesting models: {e}")
     else:
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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)}"
         )
 
 
@@ -151,7 +151,7 @@ def create_key():
                 raise e
     else:
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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)}"
         )
 
 

litellm/proxy/proxy_server.py

@@ -2305,6 +2305,94 @@ async def info_key_fn(
         )
 
 
+@router.get(
+    "/spend/keys",
+    tags=["Budget & Spend Tracking"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def spend_key_fn():
+    """
+    View all keys created, ordered by spend
+
+    Example Request: 
+    ```
+    curl -X GET "http://0.0.0.0:8000/spend/keys" \
+-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"
+            )
+
+        key_info = await prisma_client.get_data(table_name="key", query_type="find_all")
+
+        return key_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"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def view_spend_logs(
+    request_id: Optional[str] = fastapi.Query(
+        default=None,
+        description="request_id to get spend logs for specific request_id. If none passed then pass spend logs for all requests",
+    ),
+):
+    """
+    View all spend logs, if request_id is provided, only logs for that request_id will be returned
+
+    Example Request for all logs
+    ```
+    curl -X GET "http://0.0.0.0:8000/spend/logs" \
+-H "Authorization: Bearer sk-1234"
+    ```
+
+    Example Request for specific request_id
+    ```
+    curl -X GET "http://0.0.0.0:8000/spend/logs?request_id=chatcmpl-6dcb2540-d3d7-4e49-bb27-291f863f112e" \
+-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"
+            )
+        spend_logs = []
+        if request_id is not None:
+            spend_log = await prisma_client.get_data(
+                table_name="spend",
+                query_type="find_unique",
+                request_id=request_id,
+            )
+            return [spend_log]
+        else:
+            spend_logs = await prisma_client.get_data(
+                table_name="spend", query_type="find_all"
+            )
+            return spend_logs
+
+        return None
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={"error": str(e)},
+        )
+
+
 #### USER MANAGEMENT ####
 @router.post(
     "/user/new",

litellm/proxy/utils.py

@@ -361,7 +361,8 @@ class PrismaClient:
         self,
         token: Optional[str] = None,
         user_id: Optional[str] = None,
-        table_name: Optional[Literal["user", "key", "config"]] = None,
+        request_id: Optional[str] = None,
+        table_name: Optional[Literal["user", "key", "config", "spend"]] = None,
         query_type: Literal["find_unique", "find_all"] = "find_unique",
     ):
         try:
@@ -391,6 +392,10 @@ class PrismaClient:
                         for r in response:
                             if isinstance(r.expires, datetime):
                                 r.expires = r.expires.isoformat()
+                elif query_type == "find_all":
+                    response = await self.db.litellm_verificationtoken.find_many(
+                        order={"spend": "desc"},
+                    )
                 print_verbose(f"PrismaClient: response={response}")
                 if response is not None:
                     return response
@@ -407,6 +412,23 @@ class PrismaClient:
                     }
                 )
                 return response
+            elif table_name == "spend":
+                verbose_proxy_logger.debug(
+                    f"PrismaClient: get_data: table_name == 'spend'"
+                )
+                if request_id is not None:
+                    response = await self.db.litellm_spendlogs.find_unique(  # type: ignore
+                        where={
+                            "request_id": request_id,
+                        }
+                    )
+                    return response
+                else:
+                    response = await self.db.litellm_spendlogs.find_many(  # type: ignore
+                        order={"startTime": "desc"},
+                    )
+                    return response
+
         except Exception as e:
             print_verbose(f"LiteLLM Prisma Client Exception: {e}")
             import traceback

ui/admin.py

@@ -6,6 +6,9 @@ from dotenv import load_dotenv
 load_dotenv()
 import streamlit as st
 import base64, os, json, uuid, requests
+import pandas as pd
+import plotly.express as px
+import click
 
 # Replace your_base_url with the actual URL where the proxy auth app is hosted
 your_base_url = os.getenv("BASE_URL")  # Example base URL
@@ -75,7 +78,7 @@ def add_new_model():
         and st.session_state.get("proxy_key", None) is None
     ):
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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)}"
         )
 
     model_name = st.text_input(
@@ -174,10 +177,70 @@ def list_models():
             st.error(f"An error occurred while requesting models: {e}")
     else:
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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 spend_per_key():
+    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/keys"
+            else:
+                complete_url = f"{st.session_state['api_url']}/spend/keys"
+            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_key = response.json()
+                # Create DataFrame
+                spend_df = pd.DataFrame(spend_per_key)
+
+                # Display the spend per key as a graph
+                st.header("Spend ($) per API Key:")
+                top_10_df = spend_df.nlargest(10, "spend")
+                fig = px.bar(
+                    top_10_df,
+                    x="token",
+                    y="spend",
+                    title="Top 10 Spend per Key",
+                    height=550,  # Adjust the height
+                    width=1200,  # Adjust the width)
+                    hover_data=["token", "spend", "user_id", "team_id"],
+                )
+                st.plotly_chart(fig)
+
+                # Display the spend per key as a table
+                st.write("Spend per Key - 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 spend_per_user():
+    pass
+
+
 def create_key():
     import streamlit as st
     import json, requests, uuid
@@ -187,7 +250,7 @@ def create_key():
         and st.session_state.get("proxy_key", None) is None
     ):
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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)}"
         )
 
     duration = st.text_input("Duration - Can be in (h,m,s)", placeholder="1h")
@@ -235,7 +298,7 @@ def update_config():
         and st.session_state.get("proxy_key", None) is None
     ):
         st.warning(
-            "Please configure the Proxy Endpoint and Proxy Key on the Proxy Setup page."
+            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)}"
         )
 
     st.markdown("#### Alerting")
@@ -324,19 +387,25 @@ def update_config():
             raise e
 
 
-def admin_page(is_admin="NOT_GIVEN"):
+def admin_page(is_admin="NOT_GIVEN", input_api_url=None, input_proxy_key=None):
     # Display the form for the admin to set the proxy URL and allowed email subdomain
+    st.set_page_config(
+        layout="wide",  # Use "wide" layout for more space
+    )
     st.header("Admin Configuration")
     st.session_state.setdefault("is_admin", is_admin)
     # Add a navigation sidebar
     st.sidebar.title("Navigation")
+
     page = st.sidebar.radio(
         "Go to",
         (
             "Connect to Proxy",
+            "View Spend Per Key",
+            "View Spend Per User",
+            "List Models",
             "Update Config",
             "Add Models",
-            "List Models",
             "Create Key",
             "End-User Auth",
         ),
@@ -344,16 +413,23 @@ def admin_page(is_admin="NOT_GIVEN"):
     # Display different pages based on navigation selection
     if page == "Connect to Proxy":
         # Use text inputs with intermediary variables
-        input_api_url = st.text_input(
-            "Proxy Endpoint",
-            value=st.session_state.get("api_url", ""),
-            placeholder="http://0.0.0.0:8000",
-        )
-        input_proxy_key = st.text_input(
-            "Proxy Key",
-            value=st.session_state.get("proxy_key", ""),
-            placeholder="sk-...",
-        )
+        if input_api_url is None:
+            input_api_url = st.text_input(
+                "Proxy Endpoint",
+                value=st.session_state.get("api_url", ""),
+                placeholder="http://0.0.0.0:8000",
+            )
+        else:
+            st.session_state["api_url"] = input_api_url
+
+        if input_proxy_key is None:
+            input_proxy_key = st.text_input(
+                "Proxy Key",
+                value=st.session_state.get("proxy_key", ""),
+                placeholder="sk-...",
+            )
+        else:
+            st.session_state["proxy_key"] = input_proxy_key
         # When the "Save" button is clicked, update the session state
         if st.button("Save"):
             st.session_state["api_url"] = input_api_url
@@ -369,6 +445,21 @@ def admin_page(is_admin="NOT_GIVEN"):
         list_models()
     elif page == "Create Key":
         create_key()
+    elif page == "View Spend Per Key":
+        spend_per_key()
+    elif page == "View Spend Per User":
+        spend_per_user()
 
 
-admin_page()
+# admin_page()
+
+
+@click.command()
+@click.option("--proxy_endpoint", type=str, help="Proxy Endpoint")
+@click.option("--proxy_master_key", type=str, help="Proxy Master Key")
+def main(proxy_endpoint, proxy_master_key):
+    admin_page(input_api_url=proxy_endpoint, input_proxy_key=proxy_master_key)
+
+
+if __name__ == "__main__":
+    main()