Update main.py

main.py

@@ -60,6 +60,45 @@ async def verify_authorization(authorization: str = Header(None)):
     return token
 
 
+def get_gemini_models(api_key):
+    base_url = "https://generativelanguage.googleapis.com/v1beta"
+    url = f"{base_url}/models?key={api_key}"
+    
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            gemini_models = response.json()
+            return convert_to_openai_format(gemini_models)
+        else:
+            print(f"Error: {response.status_code}")
+            print(response.text)
+            return None
+    
+    except requests.RequestException as e:
+        print(f"Request failed: {e}")
+        return None
+
+def convert_to_openai_format(gemini_models):
+    openai_format = {
+        "object": "list",
+        "data": []
+    }
+    
+    for model in gemini_models.get('models', []):
+        openai_model = {
+            "id": model['name'].split('/')[-1],  # 取最后一部分作为ID
+            "object": "model",
+            "created": int(datetime.now(timezone.utc).timestamp()),  # 使用当前时间戳
+            "owned_by": "google",  # 假设所有Gemini模型都由Google拥有
+            "permission": [],  # Gemini API可能没有直接对应的权限信息
+            "root": model['name'],
+            "parent": None,  # Gemini API可能没有直接对应的父模型信息
+        }
+        openai_format["data"].append(openai_model)
+    
+    return openai_format
+        
+
 @app.get("/v1/models")
 @app.get("/hf/v1/models")
 async def list_models(authorization: str = Header(None)):
@@ -68,8 +107,7 @@ async def list_models(authorization: str = Header(None)):
         api_key = next(key_cycle)
         logger.info(f"Using API key: {api_key[:8]}...")
     try:
-        client = openai.OpenAI(api_key=api_key, base_url=config.settings.BASE_URL)
-        response = client.models.list()
+        response = get_gemini_models(api_key)
         logger.info("Successfully retrieved models list")
         return response
     except Exception as e: