From ca72ebd2dff586e7585016cd2f5b2602b10a5cfd Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 9 Sep 2023 19:11:33 -0700
Subject: [PATCH] store llm costs in budget manager

---
 litellm/budget_manager.py | 18 +++++++++++++++++-
 litellm/cost.json         |  5 +++++
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 litellm/cost.json

diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py
index fa7713156..bdda4c172 100644
--- a/litellm/budget_manager.py
+++ b/litellm/budget_manager.py
@@ -21,8 +21,24 @@ class BudgetManager:
 
     def update_cost(self, completion_obj: ModelResponse, user: str):
         cost = litellm.completion_cost(completion_response=completion_obj)
+        model = completion_obj.get("model", "") # if this throws an error try, model = completion_obj['model']
+        self.store_llm_costs(cost, model)
         self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0)
         return self.user_dict[user]["current_cost"]
     
     def get_current_cost(self, user):
-        return self.user_dict[user].get("current_cost", 0)
\ No newline at end of file
+        return self.user_dict[user].get("current_cost", 0)
+
+    def store_llm_costs(self, cost, model):
+        try:
+            import json
+            with open("cost.json", 'r') as json_file:
+                cost_dict = json.load(json_file)
+                if model in cost_dict:
+                    cost_dict[model] += cost
+                else:
+                    cost_dict[model] = cost
+                with open("cost.json", 'w') as json_file:
+                    json.dump(cost_dict, json_file, indent=4)  # Indent for pretty formatting
+        except Exception as e:
+            return # [Non blocking]
\ No newline at end of file
diff --git a/litellm/cost.json b/litellm/cost.json
new file mode 100644
index 000000000..360f981e4
--- /dev/null
+++ b/litellm/cost.json
@@ -0,0 +1,5 @@
+{
+    "gpt-3.5-turbo-0613": 0.00015000000000000001,
+    "claude-2": 0.00016454,
+    "gpt-4-0613": 0.015408
+}
\ No newline at end of file