Update api_usage.py

api_usage.py

@@ -512,8 +512,9 @@ def check_key_aws_availability(key):
     quarantine = False
     iam_full_access = False
     iam_policies_perm = False
-    enable_region_v2 = 'None'
-    enable_region_sonnet = 'None'
+    iam_user_change_password = False
+    aws_bedrock_full_access = False
+    
     session = boto3.Session(
         aws_access_key_id=access_id,
         aws_secret_access_key=access_secret
@@ -524,7 +525,7 @@ def check_key_aws_availability(key):
     username = check_username(session)
     #print(username)
     if not username[0]:
-        return False, "", "", "", "", "", username[1], "", ""
+        return False, username[1]
     
     if username[0] == 'root' and username[2]:
         root = True
@@ -540,17 +541,15 @@ def check_key_aws_availability(key):
                     iam_full_access = True
                 if policy['PolicyName'] == 'AWSCompromisedKeyQuarantineV2':
                     quarantine = True
+                if policy['PolicyName'] == 'IAMUserChangePassword':
+                    iam_user_change_password = True
+                if policy['PolicyName'] == 'AmazonBedrockFullAccess':
+                    aws_bedrock_full_access = True
                     
     enable_region = check_bedrock_invoke(session)
-    
-    if enable_region[0]:
-        enable_region_v2 = enable_region[0]
-    if enable_region[1]:
-        enable_region_sonnet = enable_region[1]
     cost = check_aws_billing(session)
-    if enable_region[0]:
-        return True, username[0], root, admin, iam_full_access, quarantine, enable_region_v2, enable_region_sonnet, cost
-    return True, username[0], root, admin, iam_full_access, quarantine, enable_region_v2, enable_region_sonnet, cost
+    
+    return True, username[0], root, admin, quarantine, iam_full_access, iam_user_change_password, aws_bedrock_full_access, enable_region, cost
     
 def check_username(session):
     try:
@@ -592,24 +591,32 @@ def invoke_claude(session, region, modelId):
         #print(e)
         return
 
+def invoke_and_collect(session, model_name, region):
+    result = invoke_claude(session, region, f"anthropic.{model_name}")
+    if result:
+        return model_name, result
+    
 def check_bedrock_invoke(session):
-    regions = ['us-east-1', 'us-west-2', 'eu-central-1', 'ap-southeast-1', 'ap-northeast-1']
-    enable_region = []
-    enable_region_sonnet = []
-    with concurrent.futures.ThreadPoolExecutor() as executer:
-        futures = [executer.submit(invoke_claude, session, region, "anthropic.claude-v2") for region in regions]            
-        for future in concurrent.futures.as_completed(futures):
-            if future.result() and (future.result() not in enable_region):
-                enable_region.append(future.result())
-        futures = [executer.submit(invoke_claude, session, region, "anthropic.claude-v2:1") for region in regions]            
-        for future in concurrent.futures.as_completed(futures):
-            if future.result() and (future.result() not in enable_region):
-                enable_region.append(future.result())      
-        futures = [executer.submit(invoke_claude, session, region, "anthropic.claude-3-sonnet-20240229-v1:0") for region in regions]
+    regions = ['us-east-1', 'us-west-2', 'eu-central-1', 'eu-west-3', 'ap-southeast-1', 'ap-northeast-1']
+    models = {
+        "claude-v2": [],
+        "claude-3-haiku-20240307-v1:0": [],
+        "claude-3-sonnet-20240229-v1:0": [],
+        "claude-3-opus-20240229-v1:0": []
+    }
+    
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        futures = []
+        for region in regions:
+            for model in models:
+                futures.append(executor.submit(invoke_and_collect, session, model, region))
+
         for future in concurrent.futures.as_completed(futures):
-            if future.result() and (future.result() not in enable_region_sonnet):
-                enable_region_sonnet.append(future.result())      
-    return enable_region, enable_region_sonnet
+            if future.result():
+                model_name, region = future.result()
+                models[model_name].append(region)
+
+    return models
 
 def check_aws_billing(session):
     try: