From f35de78df12169c3521b21b65d97ecc32af41082 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 30 Nov 2024 19:09:00 -0800
Subject: [PATCH] refactor: add more debug statements

---
 enterprise/utils.py      | 36 ++++++++++++++++++------------------
 tests/test_spend_logs.py |  2 +-
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/enterprise/utils.py b/enterprise/utils.py
index f0af1d676..d1aaa146b 100644
--- a/enterprise/utils.py
+++ b/enterprise/utils.py
@@ -2,7 +2,9 @@
 from typing import Optional, List
 from litellm._logging import verbose_logger
 from litellm.proxy.proxy_server import PrismaClient, HTTPException
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
 import collections
+import httpx
 from datetime import datetime
 
 
@@ -114,7 +116,6 @@ async def ui_get_spend_by_tags(
 
 
 def _forecast_daily_cost(data: list):
-    import requests  # type: ignore
     from datetime import datetime, timedelta
 
     if len(data) == 0:
@@ -144,9 +145,12 @@ def _forecast_daily_cost(data: list):
         today_date.year, today_date.month % 12 + 1, 1
     ) - timedelta(days=1)
 
+    print("last day of todays month", last_day_of_todays_month)
     # Calculate the remaining days in the month
     remaining_days = (last_day_of_todays_month - last_entry_date).days
 
+    print("remaining days", remaining_days)
+
     current_spend_this_month = 0
     series = {}
     for entry in data:
@@ -176,13 +180,19 @@ def _forecast_daily_cost(data: list):
         "Content-Type": "application/json",
     }
 
-    response = requests.post(
-        url="https://trend-api-production.up.railway.app/forecast",
-        json=payload,
-        headers=headers,
-    )
-    # check the status code
-    response.raise_for_status()
+    client = HTTPHandler()
+
+    try:
+        response = client.post(
+            url="https://trend-api-production.up.railway.app/forecast",
+            json=payload,
+            headers=headers,
+        )
+    except httpx.HTTPStatusError as e:
+        raise HTTPException(
+            status_code=500,
+            detail={"error": f"Error getting forecast: {e.response.text}"},
+        )
 
     json_response = response.json()
     forecast_data = json_response["forecast"]
@@ -206,13 +216,3 @@ def _forecast_daily_cost(data: list):
         f"Predicted Spend for { today_month } 2024, ${total_predicted_spend}"
     )
     return {"response": response_data, "predicted_spend": predicted_spend}
-
-    # print(f"Date: {entry['date']}, Spend: {entry['spend']}, Response: {response.text}")
-
-
-# _forecast_daily_cost(
-#     [
-#         {"date": "2022-01-01", "spend": 100},
-
-#     ]
-# )
diff --git a/tests/test_spend_logs.py b/tests/test_spend_logs.py
index a5db51a88..585079b12 100644
--- a/tests/test_spend_logs.py
+++ b/tests/test_spend_logs.py
@@ -114,7 +114,7 @@ async def test_spend_logs():
 
 
 async def get_predict_spend_logs(session):
-    url = f"http://0.0.0.0:4000/global/predict/spend/logs"
+    url = "http://0.0.0.0:4000/global/predict/spend/logs"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
     data = {
         "data": [