Update api_usage.py

api_usage.py

@@ -1,3 +1,4 @@
+from awsLib import bedrock_model_available, bedrock_send_fake_form, send_signed_request_bedrock, get_service_cost_and_usage
 import requests
 import json
 import os
@@ -9,7 +10,6 @@ import botocore.exceptions
 import concurrent.futures
 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"]
@@ -538,11 +538,12 @@ async def check_key_aws_availability(key):
                     iam_user_change_password = True
                 if policy['PolicyName'] == 'AmazonBedrockFullAccess':
                     aws_bedrock_full_access = True
-                    
-    enable_region = await check_bedrock_claude_status(access_id, access_secret)
+    async with aiohttp.ClientSession() as async_session: 
+        enable_region = await check_bedrock_claude_status(async_session, access_id, access_secret)
+        models_billing = await check_model_billing(async_session, 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
+    return True, username[0], root, admin, quarantine, iam_full_access, iam_user_change_password, aws_bedrock_full_access, enable_region, models_billing, cost
     
 def check_username(session):
     try:
@@ -569,8 +570,10 @@ def is_model_working(form_info, model_info):
         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 "Yes"
+        if 'formData' in form_status and agreement_status == 'AVAILABLE' and entitlementAvai == 'AVAILABLE':
+            if auth_status == 'AUTHORIZED':
+                return "Yes"
+            return "Maybe"
         if agreement_status == "ERROR":
             return model_info['agreementAvailability']['errorMessage']
         return "No"
@@ -583,13 +586,16 @@ async def get_model_status(session, key, secret, region, model_name, form_info):
     model_status = is_model_working(form_info, model_info)
     if model_status == "Yes":
         return region, model_name, ""
+    elif model_status == "Maybe":
+        return region, model_name, "Maybe"
     elif model_status == "No":
         return None, model_name, ""
     else:
         return None, model_name, model_status
         
-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
+async def check_bedrock_claude_status(session, key, secret):
+    # currently these regions aren't "gated" nor having only "low context" models
+    regions = ['us-east-1', 'us-west-2', 'eu-central-1', 'eu-west-3', 'ap-northeast-1', 'ap-southeast-2']
     
     models = {
         "claude-v2": [],
@@ -598,20 +604,32 @@ async def check_bedrock_claude_status(key, secret):
         "claude-3-opus-20240229-v1:0": [],
         "claude-3-5-sonnet-20240620-v1:0": []
     }
-    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:
-                tasks.append(get_model_status(session, key, secret, region, model, form_info))
-        results = await asyncio.gather(*tasks)
-        for region, model_name, msg in results:
-            if region and model_name:
+    
+    payload = json.dumps({
+        "max_tokens": 0,
+        "messages": [{"role": "user", "content": ""}],
+        "anthropic_version": "bedrock-2023-05-31"
+    })
+    
+    tasks = []
+    form_info = await bedrock_send_fake_form(session, key, secret, "us-east-1", "")
+    for region in regions:
+        for model in models:
+            tasks.append(get_model_status(session, key, secret, region, model, form_info))
+    results = await asyncio.gather(*tasks)
+    for region, model_name, msg in results:
+        if region and model_name:
+            if msg == "Maybe":
+                invoke_info = await send_signed_request_bedrock(session, payload, f"anthropic.{model_name}", key, secret, region)
+                if 'messages.0' in invoke_info.get('message'):
+                    models[model_name].append(f'{region}: might be Unavailable if disable')
+            else:
                 models[model_name].append(region)
-            elif form_info.get('message') == "Operation not allowed" and "Operation not allowed" not in models[model_name]:
-                models[model_name].append('Operation not allowed')
-            elif msg and msg not in models[model_name]:
-                models[model_name].append(msg)
+        elif form_info.get('message') == "Operation not allowed" and "Operation not allowed" not in models[model_name]:
+            models[model_name].append('Operation not allowed')
+        elif msg and msg not in models[model_name]:
+            models[model_name].append(msg)
+            
     return models
 
 def check_aws_billing(session):
@@ -629,6 +647,27 @@ def check_aws_billing(session):
     except botocore.exceptions.ClientError as error:
         return error.response['Error']['Message']
 
+async def check_model_billing(session, key, secret):
+    services = {
+        'Claude (Amazon Bedrock Edition)': 'Claude 2',
+        'Claude 3 Haiku (Amazon Bedrock Edition)': 'Claude 3 Haiku',
+        'Claude 3 Sonnet (Amazon Bedrock Edition)': 'Claude 3 Sonnet',
+        'Claude 3 Opus (Amazon Bedrock Edition)': 'Claude 3 Opus',
+        'Claude 3.5 Sonnet (Amazon Bedrock Edition)': 'Claude 3.5 Sonnet'
+    }
+    costs = {}
+
+    cost_info = await asyncio.gather(*(get_service_cost_and_usage(session, key, secret, service) for service in services))
+    for cost_and_usage, model in cost_info:
+        USD = 0
+        try:
+            for result in cost_and_usage["ResultsByTime"]:
+                USD+=float(result["Total"]["BlendedCost"]["Amount"])
+            costs[f'{services[model]} ({cost_and_usage["ResultsByTime"][0]["Total"]["BlendedCost"]["Unit"]})'] = USD
+        except:
+            costs[services[model]] = USD
+    return costs
+
 def check_key_or_availability(key):
     url = "https://openrouter.ai/api/v1/auth/key"
     headers = {'Authorization': f'Bearer {key}'}