Support downloading configuration files from external URLs.

main.py

@@ -4,38 +4,25 @@ import httpx
 import secrets
 from contextlib import asynccontextmanager
 
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi import FastAPI, HTTPException, Depends
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
 from models import RequestModel
-from utils import config, api_keys_db, api_list, error_handling_wrapper, get_all_models, verify_api_key, post_all_models, update_config
+from utils import error_handling_wrapper, get_all_models, post_all_models, load_config
 from request import get_payload
 from response import fetch_response, fetch_response_stream
 
 from typing import List, Dict
 from urllib.parse import urlparse
-from fastapi.responses import StreamingResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # 启动时的代码
     timeout = httpx.Timeout(connect=15.0, read=10.0, write=30.0, pool=30.0)
     app.state.client = httpx.AsyncClient(timeout=timeout)
-    import os
-    import yaml
-    # 新增： 从环境变量获取配置URL并拉取配置
-    config_url = os.environ.get('CONFIG_URL')
-    if config_url:
-        try:
-            response = await app.state.client.get(config_url)
-            response.raise_for_status()
-            config_data = yaml.safe_load(response.text)
-            # 更新配置
-            global config, api_keys_db, api_list
-            config, api_keys_db, api_list = update_config(config_data)
-        except Exception as e:
-            logger.error(f"Error fetching or parsing config from {config_url}: {str(e)}")
-
+    app.state.config, app.state.api_keys_db, app.state.api_list = await load_config(app)
     yield
     # 关闭时的代码
     await app.state.client.aclose()
@@ -96,6 +83,10 @@ class ModelRequestHandler:
         self.last_provider_index = -1
 
     def get_matching_providers(self, model_name, token):
+        config = app.state.config
+        # api_keys_db = app.state.api_keys_db
+        api_list = app.state.api_list
+
         api_index = api_list.index(token)
         provider_rules = []
 
@@ -127,6 +118,10 @@ class ModelRequestHandler:
         return provider_list
 
     async def request_model(self, request: RequestModel, token: str):
+        config = app.state.config
+        # api_keys_db = app.state.api_keys_db
+        api_list = app.state.api_list
+
         model_name = request.model
         matching_providers = self.get_matching_providers(model_name, token)
         # print("matching_providers", json.dumps(matching_providers, indent=4, ensure_ascii=False))
@@ -164,6 +159,16 @@ class ModelRequestHandler:
 
 model_handler = ModelRequestHandler()
 
+# 安全性依赖
+security = HTTPBearer()
+
+def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    api_list = app.state.api_list
+    token = credentials.credentials
+    if token not in api_list:
+        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
+    return token
+
 @app.post("/v1/chat/completions")
 async def request_model(request: RequestModel, token: str = Depends(verify_api_key)):
     return await model_handler.request_model(request, token)
@@ -174,7 +179,7 @@ async def options_handler():
 
 @app.post("/v1/models")
 async def list_models(token: str = Depends(verify_api_key)):
-    models = post_all_models(token)
+    models = post_all_models(token, app.state.config, app.state.api_list)
     return JSONResponse(content={
         "object": "list",
         "data": models
@@ -182,7 +187,7 @@ async def list_models(token: str = Depends(verify_api_key)):
 
 @app.get("/v1/models")
 async def list_models():
-    models = get_all_models()
+    models = get_all_models(config=app.state.config)
     return JSONResponse(content={
         "object": "list",
         "data": models

utils.py

@@ -1,4 +1,3 @@
-import yaml
 import json
 from fastapi import HTTPException, Depends
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
@@ -22,24 +21,48 @@ def update_config(config_data):
     return config_data, api_keys_db, api_list
 
 # 读取YAML配置文件
-def load_config():
+async def load_config(app):
+    import yaml
     try:
         with open('./api.yaml', 'r') as f:
             # 判断是否为空文件
             conf = yaml.safe_load(f)
+            # conf = None
             if conf:
-                return update_config(conf)
+                config, api_keys_db, api_list = update_config(conf)
             else:
-                logger.error("配置文件 'api.yaml' 为空。请检查文件内容。")
-                return [], [], []
+                # logger.error("配置文件 'api.yaml' 为空。请检查文件内容。")
+                config, api_keys_db, api_list = [], [], []
     except FileNotFoundError:
         logger.error("配置文件 'api.yaml' 未找到。请确保文件存在于正确的位置。")
-        return [], [], []
+        config, api_keys_db, api_list = [], [], []
     except yaml.YAMLError:
         logger.error("配置文件 'api.yaml' 格式不正确。请检查 YAML 格式。")
-        return [], [], []
-
-config, api_keys_db, api_list = load_config()
+        config, api_keys_db, api_list = [], [], []
+
+    if config != []:
+        return config, api_keys_db, api_list
+
+    import os
+    # 新增： 从环境变量获取配置URL并拉取配置
+    config_url = os.environ.get('CONFIG_URL')
+    if config_url:
+        try:
+            response = await app.state.client.get(config_url)
+            # logger.info(f"Fetching config from {response.text}")
+            response.raise_for_status()
+            config_data = yaml.safe_load(response.text)
+            # 更新配置
+            # logger.info(config_data)
+            if config_data:
+                config, api_keys_db, api_list = update_config(config_data)
+            else:
+                logger.error(f"Error fetching or parsing config from {config_url}")
+                config, api_keys_db, api_list = [], [], []
+        except Exception as e:
+            logger.error(f"Error fetching or parsing config from {config_url}: {str(e)}")
+            config, api_keys_db, api_list = [], [], []
+    return config, api_keys_db, api_list
 
 def ensure_string(item):
     if isinstance(item, (bytes, bytearray)):
@@ -86,7 +109,7 @@ async def error_handling_wrapper(generator, status_code=200):
         # 处理生成器为空的情况
         return async_generator(["data: {'error': 'No data returned'}\n\n"])
 
-def post_all_models(token):
+def post_all_models(token, config, api_list):
     all_models = []
     unique_models = set()
 
@@ -141,7 +164,7 @@ def post_all_models(token):
 
     return all_models
 
-def get_all_models():
+def get_all_models(config):
     all_models = []
     unique_models = set()
 
@@ -157,13 +180,4 @@ def get_all_models():
                 }
                 all_models.append(model_info)
 
-    return all_models
-
-# 安全性依赖
-security = HTTPBearer()
-
-def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
-    token = credentials.credentials
-    if token not in api_list:
-        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
-    return token
+    return all_models