From f2abac748e72b9cb5400272bdd2bca74f035377d Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 11 Jun 2024 16:16:10 -0700
Subject: [PATCH] fix refactor management endpoint utils

---
 litellm/proxy/proxy_server.py                 |  5 +--
 litellm/proxy/utils.py                        | 30 ------------------
 litellm/proxy/utils/http_parsing_utils.py     | 31 +++++++++++++++++++
 .../{ => utils}/management_endpoint_utils.py  |  7 +++--
 4 files changed, 38 insertions(+), 35 deletions(-)
 create mode 100644 litellm/proxy/utils/http_parsing_utils.py
 rename litellm/proxy/{ => utils}/management_endpoint_utils.py (90%)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 247751f0f5..640f09c731 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -103,7 +103,6 @@ from litellm.proxy.utils import (
     hash_token,
     html_form,
     missing_keys_html_form,
-    _read_request_body,
     _is_valid_team_configs,
     _is_user_proxy_admin,
     _get_user_role,
@@ -115,6 +114,8 @@ from litellm.proxy.utils import (
     _to_ns,
     get_error_message_str,
 )
+from litellm.proxy.utils.http_parsing_utils import _read_request_body
+
 from litellm import (
     CreateBatchRequest,
     RetrieveBatchRequest,
@@ -163,7 +164,7 @@ from litellm.proxy.auth.auth_checks import (
     get_actual_routes,
     log_to_opentelemetry,
 )
-from litellm.proxy.management_endpoint_utils import management_endpoint_wrapper
+from litellm.proxy.utils.management_endpoint_utils import management_endpoint_wrapper
 from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
 from litellm.exceptions import RejectedRequestError
 from litellm.integrations.slack_alerting import SlackAlertingArgs, SlackAlerting
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 3530b93347..93790e7641 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -2596,36 +2596,6 @@ async def update_spend(
                 raise e
 
 
-async def _read_request_body(request: Optional[Request]) -> dict:
-    """
-    Asynchronous function to read the request body and parse it as JSON or literal data.
-
-    Parameters:
-    - request: The request object to read the body from
-
-    Returns:
-    - dict: Parsed request data as a dictionary
-    """
-    import ast, json
-
-    try:
-        request_data: dict = {}
-        if request is None:
-            return request_data
-        body = await request.body()
-
-        if body == b"" or body is None:
-            return request_data
-        body_str = body.decode()
-        try:
-            request_data = ast.literal_eval(body_str)
-        except:
-            request_data = json.loads(body_str)
-        return request_data
-    except:
-        return {}
-
-
 def _is_projected_spend_over_limit(
     current_spend: float, soft_budget_limit: Optional[float]
 ):
diff --git a/litellm/proxy/utils/http_parsing_utils.py b/litellm/proxy/utils/http_parsing_utils.py
new file mode 100644
index 0000000000..cc24713919
--- /dev/null
+++ b/litellm/proxy/utils/http_parsing_utils.py
@@ -0,0 +1,31 @@
+from typing import Optional
+from fastapi import Request
+import ast, json
+
+
+async def _read_request_body(request: Optional[Request]) -> dict:
+    """
+    Asynchronous function to read the request body and parse it as JSON or literal data.
+
+    Parameters:
+    - request: The request object to read the body from
+
+    Returns:
+    - dict: Parsed request data as a dictionary
+    """
+    try:
+        request_data: dict = {}
+        if request is None:
+            return request_data
+        body = await request.body()
+
+        if body == b"" or body is None:
+            return request_data
+        body_str = body.decode()
+        try:
+            request_data = ast.literal_eval(body_str)
+        except:
+            request_data = json.loads(body_str)
+        return request_data
+    except:
+        return {}
diff --git a/litellm/proxy/management_endpoint_utils.py b/litellm/proxy/utils/management_endpoint_utils.py
similarity index 90%
rename from litellm/proxy/management_endpoint_utils.py
rename to litellm/proxy/utils/management_endpoint_utils.py
index aabf6e5d3b..33d9969923 100644
--- a/litellm/proxy/management_endpoint_utils.py
+++ b/litellm/proxy/utils/management_endpoint_utils.py
@@ -1,7 +1,7 @@
 from datetime import datetime
 from functools import wraps
 from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload
-from litellm.proxy.utils import _read_request_body
+from http_parsing_utils import _read_request_body
 from fastapi import Request
 
 
@@ -20,8 +20,9 @@ def management_endpoint_wrapper(func):
         try:
             result = await func(*args, **kwargs)
             end_time = datetime.now()
-            user_api_key_dict: UserAPIKeyAuth = kwargs["user_api_key_dict"]
-
+            if kwargs is None:
+                kwargs = {}
+            user_api_key_dict: UserAPIKeyAuth = kwargs.get("user_api_key_dict")
             parent_otel_span = user_api_key_dict.parent_otel_span
             if parent_otel_span is not None:
                 from litellm.proxy.proxy_server import open_telemetry_logger