Merge pull request #103 from AlexHTW/add-balance-function

Add balance functions to stats command

bot/openai_helper.py

@@ -5,6 +5,11 @@ import tiktoken
 
 import openai
 
+import requests
+import json
+from datetime import date
+from calendar import monthrange
+
 # Models can be found here: https://platform.openai.com/docs/models/overview
 GPT_3_MODELS = ("gpt-3.5-turbo", "gpt-3.5-turbo-0301")
 GPT_4_MODELS = ("gpt-4", "gpt-4-0314")
@@ -266,3 +271,38 @@ class OpenAIHelper:
                     num_tokens += tokens_per_name
         num_tokens += 2  # every reply is primed with <im_start>assistant
         return num_tokens
+
+    def get_grant_balance(self):
+        """Gets remaining grant balance for new users from OpenAI API.
+
+        :return: remaining grant balance
+        """
+        headers = {
+            "Authorization": f"Bearer {openai.api_key}"
+        }
+        response = requests.get("https://api.openai.com/dashboard/billing/credit_grants", headers=headers)
+        billing_data = json.loads(response.text)
+        balance = billing_data["total_available"]
+        return balance
+    
+    def get_billing_current_month(self):
+        """Gets billed usage for current month from OpenAI API.
+
+        :return: dollar amount of usage this month
+        """
+        headers = {
+            "Authorization": f"Bearer {openai.api_key}"
+        }
+        # calculate first and last day of current month
+        today = date.today()
+        first_day = date(today.year, today.month, 1)
+        _, last_day_of_month = monthrange(today.year, today.month)
+        last_day = date(today.year, today.month, last_day_of_month)
+        params = {
+            "start_date": first_day,
+            "end_date": last_day
+        }
+        response = requests.get("https://api.openai.com/dashboard/billing/usage", headers=headers, params=params)
+        billing_data = json.loads(response.text)
+        usage_month = billing_data["total_usage"] / 100 # convert cent amount to dollars
+        return usage_month