From de631d3d918c16954cbe00b015de3170c035dcc9 Mon Sep 17 00:00:00 2001
From: reanon <85157-reanon@users.noreply.gitgud.io>
Date: Thu, 17 Apr 2025 20:13:06 +0000
Subject: [PATCH] Edit google-ai.ts

---
 src/proxy/google-ai.ts | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/proxy/google-ai.ts b/src/proxy/google-ai.ts
index 777b808..bb21374 100644
--- a/src/proxy/google-ai.ts
+++ b/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;