Update api_usage.py

api_usage.py

@@ -434,12 +434,16 @@ def check_key_aws_availability(key):
                     quarantine = True
                     
     enable_region = check_bedrock_invoke(session)
+    enable_region_v2 = "None"
+    enable_region_sonnet = "None"
+    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:
-        return True, username[0], root, admin, quarantine, enable_region, cost
-    if root or admin:
-        return True, username[0], root, admin, quarantine, "None", cost
-    return True, username[0], root, admin, quarantine, "None", cost
+    if enable_region[0]:
+        return True, username[0], root, admin, quarantine, enable_region_v2, enable_region_sonnet, cost
+    return True, username[0], root, admin, quarantine, enable_region_v2, enable_region_sonnet, cost
     
 def check_username(session):
     try:
@@ -469,13 +473,13 @@ def invoke_claude(session, region, modelId):
         response = bedrock_runtime.invoke_model(body=body, modelId=modelId)
     except bedrock_runtime.exceptions.ValidationException as error:
         #print(error.response['Error'])
-        return region
+        if 'max_tokens_to_sample' in error.response['Error']['Message']:
+            return region
     except bedrock_runtime.exceptions.AccessDeniedException as error:
         #print(error.response['Error'])
         return
     except bedrock_runtime.exceptions.ResourceNotFoundException as error:
         #print(error.response['Error'])
-        return invoke_claude(session, region, "anthropic.claude-v2:1")
     except Exception as e:
         #print(e)
         return
@@ -483,12 +487,21 @@ def invoke_claude(session, region, modelId):
 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():
-                enable_region.append(future.result())            
-    return 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():
+                enable_region.append(future.result())      
+        futures = [executer.submit(invoke_claude, session, region, "anthropic.claude-3-sonnet-20240229-v1:0") for region in regions]
+        for future in concurrent.futures.as_completed(futures):
+            if future.result():
+                enable_region_sonnet.append(future.result())      
+    return enable_region, enable_region_sonnet
 
 def check_aws_billing(session):
     try: