Update api_usage.py

api_usage.py

@@ -33,7 +33,7 @@ def get_headers(key, org_id:str = None):
         headers["OpenAI-Organization"] = org_id
     return headers
 
-def get_subscription(key, org_list):
+def get_subscription(key, session, org_list):
     has_gpt4 = False
     has_gpt4_32k = False
     default_org = ""
@@ -46,7 +46,7 @@ def get_subscription(key, org_list):
     list_models_avai = set()
     
     for org_in in org_list:
-        available_models = get_models(key, org_in['id'])
+        available_models = get_models(session, key, org_in['id'])
         headers = get_headers(key, org_in['id'])
         has_gpt4_32k = True if GPT_TYPES[2] in available_models else False
         has_gpt4 = True if GPT_TYPES[1] in available_models else False
@@ -57,7 +57,7 @@ def get_subscription(key, org_list):
         if has_gpt4_32k:
             org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
             list_models_avai.update(GPT_TYPES)
-            status_formated = format_status([GPT_TYPES[2], GPT_TYPES[1], GPT_TYPES[0]], headers)
+            status_formated = format_status([GPT_TYPES[2], GPT_TYPES[1], GPT_TYPES[0]], session, headers)
             rpm.append(status_formated[0])
             tpm.append(status_formated[1])
             quota.append(status_formated[2])
@@ -66,7 +66,7 @@ def get_subscription(key, org_list):
         elif has_gpt4:
             org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
             list_models_avai.update([GPT_TYPES[1], GPT_TYPES[0]])
-            status_formated = format_status([GPT_TYPES[1], GPT_TYPES[0]], headers)
+            status_formated = format_status([GPT_TYPES[1], GPT_TYPES[0]], session, headers)
             rpm.append(status_formated[0])
             tpm.append(status_formated[1])
             quota.append(status_formated[2])
@@ -75,7 +75,7 @@ def get_subscription(key, org_list):
         else:
             org.append(f"{org_in['id']} ({org_in['name']}, {org_in['title']}, {org_in['role']})")
             list_models_avai.update([GPT_TYPES[0]])
-            status_formated = format_status([GPT_TYPES[0]], headers)
+            status_formated = format_status([GPT_TYPES[0]], session, headers)
             rpm.append(status_formated[0])
             tpm.append(status_formated[1])
             quota.append(status_formated[2])
@@ -139,12 +139,11 @@ def send_oai_completions(oai_stuff):
 def helper_oai(oai_stuff):
     return send_oai_completions(oai_stuff)
         
-def format_status(list_models_avai, headers):
+def format_status(list_models_avai, session, headers):
     rpm = []
     tpm = []
     quota = ""
-    r = requests.Session()
-    args = [(r, headers, model) for model in list_models_avai]
+    args = [(session, headers, model) for model in list_models_avai]
     with concurrent.futures.ThreadPoolExecutor() as executer:
         for result in executer.map(helper_oai, args):
             rpm.append(result[0])
@@ -171,23 +170,23 @@ def check_key_tier(rpm, tpm, dict, headers):
         if (dictCount == dictItemsCount):
             return "yes | custom-tier"
 
-def get_orgs(key):
+def get_orgs(session, key):
     headers=get_headers(key)
-    rq = requests.get(f"{BASE_URL}/organizations", headers=headers, timeout=10)
+    rq = session.get(f"{BASE_URL}/organizations", headers=headers, timeout=10)
     return rq.json()['data']
 
-def get_models(key, org: str = None):
+def get_models(session, key, org: str = None):
     if org != None:
         headers = get_headers(key, org)
     else:
         headers = get_headers(key)
-    rq = requests.get(f"{BASE_URL}/models", headers=headers, timeout=10)
+    rq = session.get(f"{BASE_URL}/models", headers=headers, timeout=10)
     avai_models = rq.json()
     return [model["id"] for model in avai_models["data"]] #[model["id"] for model in avai_models["data"] if model["id"] in GPT_TYPES]
         
-def check_key_availability(key):
+def check_key_availability(session, key):
     try:
-        return get_orgs(key)
+        return get_orgs(session, key)
     except Exception as e:
         return False