From 6372f929d45bf55151792bf2c0f07ff2851bee38 Mon Sep 17 00:00:00 2001
From: based <phrog@riseup.net>
Date: Fri, 16 Feb 2024 15:04:52 +1000
Subject: [PATCH] distinguish makersuite keys with special gemini models

---
 APIKey.py     |  3 +++
 MakerSuite.py | 18 ++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/APIKey.py b/APIKey.py
index 6b36174..2204780 100644
--- a/APIKey.py
+++ b/APIKey.py
@@ -45,6 +45,9 @@ class APIKey:
         elif provider == Provider.MISTRAL:
             self.subbed = False
 
+        elif provider == Provider.MAKERSUITE:
+            self.models = []
+
 
 class Provider(Enum):
     OPENAI = 1,
diff --git a/MakerSuite.py b/MakerSuite.py
index 10cf580..45cca0c 100644
--- a/MakerSuite.py
+++ b/MakerSuite.py
@@ -1,16 +1,30 @@
 import APIKey
 
+gemini_models = ["gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.0-ultra"]
+
 
 async def check_makersuite(key: APIKey, session):
     async with session.get(f"https://generativelanguage.googleapis.com/v1beta/models?key={key.api_key}") as response:
         if response.status != 200:
             return
+        response_json = await response.json()
+        model_names = [model['name'].replace('models/', '') for model in response_json['models']]
+        for model in gemini_models:
+            if model in model_names:
+                key.models.append(model)
         return True
 
 
 def pretty_print_makersuite_keys(keys):
+    model_counts = {model: [] for model in gemini_models}
+    for key in keys:
+        for model in key.models:
+            model_counts[model].append(key.api_key)
+
     print('-' * 90)
     print(f'Validated {len(keys)} MakerSuite keys:')
-    for key in keys:
-        print(f'{key.api_key}')
+    for model, keys in model_counts.items():
+        print(f'\n{len(keys)} keys with model {model}:')
+        for key in keys:
+            print(f'{key}')
     print(f'\n--- Total Valid MakerSuite Keys: {len(keys)} ---\n')