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); + } + } + nativeModelsCache = response.data; nativeModelsCacheTime = new Date().getTime(); return nativeModelsCache;