Update app.py

app.py

@@ -12,7 +12,7 @@ from api_usage import (
     get_orgs_me,
     check_key_ant_availability,
     check_ant_rate_limit,
-    check_key_gemini_availability, # 이 함수는 (bool, str) 튜플을 반환한다고 가정
+    check_key_gemini_availability,
     check_key_azure_availability,
     get_azure_status,
     get_azure_deploy,
@@ -204,19 +204,22 @@ async def get_key_ant_info(key: str, rate_limit: bool, claude_model: str) -> Dic
 
     return info_dict
 
-
-# 사용자의 원본 get_key_gemini_info 함수 유지
+########## 변경된 부분 ##########
+# 1. get_key_gemini_info 함수가 모델 목록을 포함하도록 수정
 def get_key_gemini_info(key: str) -> Dict[str, Any]:
-    """Gemini 키 정보를 가져옵니다 (사용자 원본 버전)."""
-    key_avai = check_key_gemini_availability(key) # (bool, str) 튜플 반환 가정
+    """Gemini 키 정보를 가져오고 모델 목록을 포함합니다."""
+    # check_key_gemini_availability가 (bool, str, list_of_models) 튜플을 반환한다고 가정
+    key_avai = check_key_gemini_availability(key)
     info_dict = {
-        "key": key, # 원본에는 key 필드가 있었음
+        "key": key,
+        "key_type": "Google Gemini",
         "key_availability": key_avai[0],
         "status": key_avai[1],
+        # 반환값 길이를 확인하여 안전하게 모델 목록을 가져옴
+        "models": key_avai[2] if len(key_avai) > 2 else []
     }
     return info_dict
 
-
 def get_key_azure_info(endpoint: str, api_key: str) -> Dict[str, Any]:
     key_avai = check_key_azure_availability(endpoint, api_key)
     info_dict = {
@@ -485,15 +488,24 @@ async def process_single_key(key: str, rate_limit: bool, claude_model: str) -> D
     if _key.startswith("sk-"):
         result = get_key_oai_info(_key)
         return {"key": _key, **result}
+
+    ########## 변경된 부분 ##########
+    # 2. process_single_key 함수에서 특정 모델 존재 여부 확인 로직 추가
     if _key.startswith("AIzaSy"):
         gemini_info = get_key_gemini_info(_key)
         is_working = gemini_info.get("key_availability") and gemini_info.get("status") == "Working"
+
+        # 특정 모델("gemini-2.5-pro-preview-06-05")이 사용 가능한 모델 목록에 있는지 확인
+        available_models = gemini_info.get("models", [])
+        has_target_model = "gemini-2.5-pro-preview-06-05" in available_models
+
         result = {
-            "key_type": "Google Gemini",
             **gemini_info,
-            "is_gemini_working": is_working
+            "is_gemini_working": is_working,
+            "has_specific_gemini_model": has_target_model # 결과를 저장할 플래그 추가
         }
         return result
+
     if _key.startswith("pst-"):
         result = get_key_nai_info(_key)
         return {"key": _key, **result}
@@ -562,7 +574,10 @@ async def sort_keys(text: str, rate_limit: bool, claude_model: str) -> Tuple[Lis
 
         key_type = result.get("key_type")
 
-        if result.get("is_gemini_working"):
+        ########## 변경된 부분 ##########
+        # 3. sort_keys 함수에서 Gemini 키 필터링 조건 강화
+        # 'is_gemini_working'과 'has_specific_gemini_model' 플래그를 모두 확인
+        if result.get("is_gemini_working") and result.get("has_specific_gemini_model"):
             working_gemini_keys.append(key_value)
         elif key_type == "OpenAI" and \
              result.get("key_availability") is True and \
@@ -645,9 +660,12 @@ This version highlights working OpenAI (with sufficient quota), Anthropic, and G
                 max_lines=10,
                 interactive=False,
             )
+
+            ########## 변경된 부분 ##########
+            # 4. Gemini 결과 상자의 레이블과 설명을 업데이트하여 새로운 필터링 기준을 명시
             gemini_keys_output = gr.Textbox(
-                label="Working Gemini Keys",
-                info="Lists Gemini keys confirmed as 'Working'.",
+                label="Working Gemini Keys (with gemini-2.5-pro-preview-06-05)",
+                info="Lists Gemini keys that are 'Working' AND have 'gemini-2.5-pro-preview-06-05' in their available models.",
                 lines=3,
                 max_lines=10,
                 interactive=False,
@@ -666,5 +684,4 @@ This version highlights working OpenAI (with sufficient quota), Anthropic, and G
         api_name="sort_keys",
     )
 
-# demo.launch(share=True)
 demo.launch()