Update api_usage.py

api_usage.py

@@ -7,8 +7,9 @@ from dateutil.relativedelta import relativedelta
 import boto3
 import botocore.exceptions
 import concurrent.futures
-import asyncio
+import asyncio, aiohttp
 import aiohttp
+from awsLib import bedrock_model_available,bedrock_send_fake_form
 
 BASE_URL = 'https://api.openai.com/v1'
 GPT_TYPES = ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "gpt-4-32k-0314", "gpt-4o", "gpt-4-turbo"]
@@ -20,13 +21,15 @@ TOKEN_LIMIT_PER_TIER_TURBO = {
     "tier-2": 80000,
     "tier-3": 160000,
     "tier-4": 1000000,
-    "tier-5": 2000000
+    "tier-5-old": 2000000,
+    "tier-5": 5000000
 }
 TOKEN_LIMIT_PER_TIER_GPT4 = {
     "tier-1": 10000,
     "tier-2": 40000,
     "tier-3": 80000,
-    "tier-4-5": 300000
+    "tier-4": 300000,
+    "tier-5": 1000000
 } # according to: https://platform.openai.com/docs/guides/rate-limits/usage-tiers
 
 RPM_LIMIT_PER_BUILD_TIER_ANT = {
@@ -340,9 +343,9 @@ def check_key_gemini_availability(key):
 def check_key_azure_availability(endpoint, api_key):
     try:        
         if endpoint.startswith('http'):
-            url = f'{endpoint}/openai/models?api-version=2023-03-15-preview'
+            url = f'{endpoint}/openai/models?api-version=2022-12-01'
         else:
-            url = f'https://{endpoint}/openai/models?api-version=2023-03-15-preview'
+            url = f'https://{endpoint}/openai/models?api-version=2022-12-01'
         
         headers = {
             'User-Agent': 'OpenAI/v1 PythonBindings/0.28.0', 
@@ -359,9 +362,9 @@ def check_key_azure_availability(endpoint, api_key):
 def get_azure_deploy(endpoint, api_key):
     try:        
         if endpoint.startswith('http'):
-            url = f'{endpoint}/openai/deployments?api-version=2023-03-15-preview'
+            url = f'{endpoint}/openai/deployments?api-version=2022-12-01'
         else:
-            url = f'https://{endpoint}/openai/deployments?api-version=2023-03-15-preview'
+            url = f'https://{endpoint}/openai/deployments?api-version=2022-12-01'
             
         headers = {
 			'User-Agent': 'OpenAI/v1 PythonBindings/0.28.0', 
@@ -379,9 +382,9 @@ def get_azure_deploy(endpoint, api_key):
 def check_gpt4turbo(endpoint, api_key, deploy_id):
     try:        
         if endpoint.startswith('http'):
-            url = f'{endpoint}/openai/deployments/{deploy_id}/chat/completions?api-version=2023-03-15-preview'
+            url = f'{endpoint}/openai/deployments/{deploy_id}/chat/completions?api-version=2024-02-01'
         else:
-            url = f'https://{endpoint}/openai/deployments/{deploy_id}/chat/completions?api-version=2023-03-15-preview'
+            url = f'https://{endpoint}/openai/deployments/{deploy_id}/chat/completions?api-version=2024-02-01'
             
         headers = {
             'Content-Type': 'application/json',
@@ -442,9 +445,9 @@ def get_azure_status(endpoint, api_key, deployments_list):
         
         for model, deployment in list_model.items():
             if endpoint.startswith('http'):
-                url = f'{endpoint}/openai/deployments/{deployment}/chat/completions?api-version=2023-03-15-preview'
+                url = f'{endpoint}/openai/deployments/{deployment}/chat/completions?api-version=2024-02-01'
             else:
-                url = f'https://{endpoint}/openai/deployments/{deployment}/chat/completions?api-version=2023-03-15-preview'
+                url = f'https://{endpoint}/openai/deployments/{deployment}/chat/completions?api-version=2024-02-01'
                 
             headers = {
                 'Content-Type': 'application/json',
@@ -525,7 +528,7 @@ def check_key_replicate_availability(key):
     except:
         return "Unknown", "", "", "Error while making request"
 
-def check_key_aws_availability(key):
+async def check_key_aws_availability(key):
     access_id = key.split(':')[0]
     access_secret = key.split(':')[1]
     
@@ -569,7 +572,7 @@ def check_key_aws_availability(key):
                 if policy['PolicyName'] == 'AmazonBedrockFullAccess':
                     aws_bedrock_full_access = True
                     
-    enable_region = check_bedrock_invoke(session)
+    enable_region = await check_bedrock_claude_status(access_id, access_secret)
     cost = check_aws_billing(session)
     
     return True, username[0], root, admin, quarantine, iam_full_access, iam_user_change_password, aws_bedrock_full_access, enable_region, cost
@@ -592,51 +595,47 @@ def check_policy(iam, username):
     except botocore.exceptions.ClientError as error:
         return False, error.response['Error']['Code']
 
-def invoke_claude(session, region, modelId):
+def is_model_working(form_info, model_info):
     try:
-        bedrock_runtime = session.client("bedrock-runtime", region_name=region)
-        body = json.dumps({
-            "prompt": "\n\nHuman:\n\nAssistant:",
-            "max_tokens_to_sample": 0
-        })
-        response = bedrock_runtime.invoke_model(body=body, modelId=modelId)
-    except bedrock_runtime.exceptions.ValidationException as error:
-        #print(error.response['Error'])
-        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
-    except Exception as e:
-        #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
+        form_status = form_info['message']
+        agreement_status = model_info['agreementAvailability']['status']
+        auth_status = model_info['authorizationStatus']
+        entitlementAvai = model_info['entitlementAvailability']
+        
+        if 'formData' in form_status and agreement_status == 'AVAILABLE' and auth_status == 'AUTHORIZED' and entitlementAvai == 'AVAILABLE':
+            return True
+        return False
+    except:
+        #print(form_status)
+        return False
+        
+async def get_model_status(session, key, secret, region, model_name, form_info):
+    model_info = await bedrock_model_available(session, key, secret, region, f"anthropic.{model_name}")
+    model_status = is_model_working(form_info, model_info)
+    if model_status:
+        return region, model_name
+    else:
+        return None, None
+        
+async def check_bedrock_claude_status(key, secret):
+    regions = ['us-east-1', 'us-west-2', 'eu-central-1', 'eu-west-3', 'ap-northeast-1', 'ap-southeast-2'] # currently these regions aren't "gated" nor having only "low context" models
     
-def check_bedrock_invoke(session):
-    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": []
+        "claude-3-opus-20240229-v1:0": [],
+        "claude-3-5-sonnet-20240620-v1:0": []
     }
-    
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        futures = []
+    async with aiohttp.ClientSession() as session:
+        tasks = []
+        form_info = await bedrock_send_fake_form(session, key, secret, "us-east-1", "")
         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():
-                model_name, region = future.result()
+                tasks.append(get_model_status(session, key, secret, region, model, form_info))
+        results = await asyncio.gather(*tasks)
+        for region, model_name in results:
+            if region and model_name:
                 models[model_name].append(region)
 
     return models