Edit google-ai.ts

src/proxy/google-ai.ts

@@ -38,9 +38,16 @@ const getModelsResponse = () => {
     return modelsCache;
   }
 
+  // Get all model IDs from keys, including flash and pro models
   const modelIds = Array.from(
     new Set(keys.map((k) => k.modelIds).flat())
-  ).filter((id) => id.startsWith("models/gemini"));
+  ).filter((id) => id.startsWith("models/gemini") || 
+                   id.startsWith("models/") && (
+                     id.includes("flash") || 
+                     id.includes("pro") || 
+                     id.includes("vision")
+                   ));
+                   
   const models = modelIds.map((id) => ({
     id,
     object: "model",
@@ -89,6 +96,15 @@ const getNativeModelsResponse = async () => {
       },
     });
 
+    // Update modelIds for all keys to ensure they reflect the latest models
+    const fetchedModelIds = response.data.models?.map((model: any) => model.name) || [];
+    for (const k of keys) {
+      if (!k.isDisabled && !k.isRevoked) {
+        // Update this key's modelIds with the fetched models
+        keyPool.update(k.hash, { modelIds: fetchedModelIds } as Partial<GoogleAIKey>);
+      }
+    }
+
     nativeModelsCache = response.data;
     nativeModelsCacheTime = new Date().getTime();
     return nativeModelsCache;