From 5ac1a6a93cba10a2c2bbe7b6e6ff04b168e973db Mon Sep 17 00:00:00 2001
From: based <phrog@riseup.net>
Date: Sat, 16 Mar 2024 22:15:28 +1000
Subject: [PATCH] claude quota check + aws v3 fallback for keys without v2
 enabled

---
 APIKey.py    |  1 +
 AWS.py       |  9 ++++++---
 Anthropic.py | 24 ++++++++++++++++--------
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/APIKey.py b/APIKey.py
index f144188..9f462b9 100644
--- a/APIKey.py
+++ b/APIKey.py
@@ -18,6 +18,7 @@ class APIKey:
         elif provider == Provider.ANTHROPIC:
             self.pozzed = False
             self.rate_limited = False
+            self.has_quota = True
 
         elif provider == Provider.AI21:
             self.trial_elapsed = False
diff --git a/AWS.py b/AWS.py
index dfdb889..a6f6f10 100644
--- a/AWS.py
+++ b/AWS.py
@@ -23,6 +23,10 @@ def check_aws(key: APIKey):
             session = boto3.Session(aws_access_key_id=access_key, aws_secret_access_key=secret)
             test_invoke_perms(session, key)
 
+            # retest with v3 in the case goons didn't enable v2
+            if not key.bedrock_enabled:
+                test_invoke_perms(session, key, "anthropic.claude-3-sonnet-20240229-v1:0")
+
             if key.useless:
                 key.useless_reasons.append('Failed Invoke Check')
 
@@ -50,7 +54,7 @@ def check_aws(key: APIKey):
         return
 
 
-def test_invoke_perms(session, key: APIKey):
+def test_invoke_perms(session, key: APIKey, model="anthropic.claude-v2"):
     data = {
         "prompt": "\n\nHuman:\n\nAssistant:",
         "max_tokens_to_sample": -1,
@@ -59,8 +63,7 @@ def test_invoke_perms(session, key: APIKey):
         bedrock_runtime_client = None
         try:
             bedrock_runtime_client = session.client("bedrock-runtime", region_name=region)
-            # v1 is not enabled in some regions where v2 is, proxies will report keys in these regions as logged and will by default ignore them despite them working fine for v2
-            bedrock_runtime_client.invoke_model(body=json.dumps(data), modelId="anthropic.claude-v2")
+            bedrock_runtime_client.invoke_model(body=json.dumps(data), modelId=model)
         except bedrock_runtime_client.exceptions.ValidationException as e:
             if 'max_tokens_to_sample' in e.response['Error']['Message']:
                 if key.region == "":
diff --git a/Anthropic.py b/Anthropic.py
index 0073911..10427cc 100644
--- a/Anthropic.py
+++ b/Anthropic.py
@@ -23,6 +23,10 @@ async def check_anthropic(key: APIKey, session):
         text = await response.text()
         if "This organization has been disabled" in text:
             return
+        elif "Your credit balance is too low to access the Claude API" in text:
+            key.has_quota = False
+            return True
+
         key.pozzed = any(message in text for message in pozzed_messages)
 
         return True
@@ -30,13 +34,17 @@ async def check_anthropic(key: APIKey, session):
 
 def pretty_print_anthropic_keys(keys):
     print('-' * 90)
-    pozzed = 0
-    rate_limited = 0
     print(f'Validated {len(keys)} working Anthropic keys:')
-    for key in keys:
-        if key.pozzed:
-            pozzed += 1
-        elif key.rate_limited:
-            rate_limited += 1
+    keys_with_quota = [key for key in keys if key.has_quota]
+    keys_without_quota = [key for key in keys if not key.has_quota]
+
+    pozzed = sum(key.pozzed for key in keys_with_quota)
+    rate_limited = sum(key.rate_limited for key in keys_with_quota)
+
+    print(f'Total keys with quota: {len(keys_with_quota)} (pozzed: {pozzed}, unpozzed: {len(keys_with_quota) - pozzed - rate_limited}, unsure/rate limited: {rate_limited})')
+    for key in keys_with_quota:
         print(f'{key.api_key}' + (' | pozzed' if key.pozzed else "") + (' | rate limited' if key.rate_limited else ""))
-    print(f'\n--- Total Valid Anthropic Keys: {len(keys)} ({pozzed} pozzed, {len(keys) - pozzed - rate_limited} unpozzed, {rate_limited} unsure/rate limited) ---\n')
+
+    print(f'\nTotal keys without quota: {len(keys_without_quota)}')
+    for key in keys_without_quota:
+        print(f'{key.api_key}')