Update api_usage.py

api_usage.py

@@ -476,7 +476,7 @@ def check_key_aws_availability(key):
     if not username[0]:
         return False, "", "", "", "", "", username[1], "", ""
     
-    if username[0] == 'root':
+    if username[0] == 'root' and username[2]:
         root = True
         admin = True
             
@@ -507,9 +507,9 @@ def check_username(session):
         sts = session.client('sts')
         sts_iden = sts.get_caller_identity()
         if len(sts_iden['Arn'].split('/')) > 1:
-            return sts_iden['Arn'].split('/')[1], "Valid"
+            return sts_iden['Arn'].split('/')[1], "Valid", False
             
-        return sts_iden['Arn'].split(':')[5], "Valid"
+        return sts_iden['Arn'].split(':')[5], "Valid", True
     except botocore.exceptions.ClientError as error:
         return False, error.response['Error']['Code']
         
@@ -575,7 +575,48 @@ def check_aws_billing(session):
         return ce_cost['ResultsByTime']
     except botocore.exceptions.ClientError as error:
         return error.response['Error']['Message']
-  
+
+def check_key_or_availability(key):
+    url = "https://openrouter.ai/api/v1/auth/key"
+    headers = {'Authorization': f'Bearer {key}'}
+    
+    rq = requests.get(url, headers=headers)
+    res = rq.json()
+    if rq.status_code == 200:
+        data = res['data']
+        rpm = data['rate_limit']['requests'] // int(data['rate_limit']['interval'].replace('s', '')) * 60
+        return True, data, rpm
+    return False, f"{res['error']['code']}: {res['error']['message']}", 0
+
+def check_key_or_limits(key):
+    url = "https://openrouter.ai/api/v1/models"
+    headers = {"Authorization": f"Bearer {key}"}
+    models = {
+        "openai/gpt-4-turbo-preview": "",
+        "anthropic/claude-3-sonnet:beta": "",
+        "anthropic/claude-3-opus:beta":""
+    }
+    
+    rq = requests.get(url, headers=headers)
+    res = rq.json()
+    
+    balance = 0.0
+    count = 0
+
+    for model in res['data']:
+        if model['id'] in models.keys():
+            if count == 3:
+                break
+            prompt_tokens_limit = int(model.get("per_request_limits", "").get("prompt_tokens", ""))
+            completion_tokens_limit = int(model.get("per_request_limits", "").get("completion_tokens", ""))
+            models[model['id']] = { "Prompt": prompt_tokens_limit, "Completion": completion_tokens_limit }
+            if model['id'] == "anthropic/claude-3-sonnet:beta":
+                price_prompt = float(model.get("pricing", 0).get("prompt", 0))
+                price_completion = float(model.get("pricing", 0).get("completion", 0))
+                balance = (prompt_tokens_limit * price_prompt) + (completion_tokens_limit * price_completion)
+            count+=1
+    return balance, models
+    
 if __name__ == "__main__":
     key = os.getenv("OPENAI_API_KEY")
     key_ant = os.getenv("ANTHROPIC_API_KEY")