Refactor code

main.py

@@ -1,21 +1,20 @@
 import json
-import yaml
+import traceback
+
 import httpx
-import logging
 import secrets
-import traceback
 from contextlib import asynccontextmanager
 
-from fastapi import FastAPI, Request, HTTPException, Depends
-from fastapi.responses import StreamingResponse
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from fastapi import FastAPI, HTTPException, Depends
 
 from models import RequestModel
+from utils import config, api_keys_db, api_list, error_handling_wrapper, get_all_models, verify_api_key
 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
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -28,66 +27,6 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-# 安全性依赖
-security = HTTPBearer()
-
-# 读取YAML配置文件
-def load_config():
-    try:
-        with open('./api.yaml', 'r') as f:
-            conf = yaml.safe_load(f)
-            for index, provider in enumerate(conf['providers']):
-                model_dict = {}
-                for model in provider['model']:
-                    if type(model) == str:
-                        model_dict[model] = model
-                    if type(model) == dict:
-                        model_dict.update({value: key for key, value in model.items()})
-                provider['model'] = model_dict
-                conf['providers'][index] = provider
-            api_keys_db = conf['api_keys']
-            api_list = [item["api"] for item in api_keys_db]
-            # print(json.dumps(conf, indent=4, ensure_ascii=False))
-            return conf, api_keys_db, api_list
-    except FileNotFoundError:
-        print("配置文件 'api.yaml' 未找到。请确保文件存在于正确的位置。")
-        return [], [], []
-    except yaml.YAMLError:
-        print("配置文件 'api.yaml' 格式不正确。请检查 YAML 格式。")
-        return [], [], []
-
-config, api_keys_db, api_list = load_config()
-
-async def error_handling_wrapper(generator, status_code=200):
-    try:
-        first_item = await generator.__anext__()
-        first_item_str = first_item
-        if isinstance(first_item_str, (bytes, bytearray)):
-            first_item_str = first_item_str.decode("utf-8")
-        if isinstance(first_item_str, str):
-            if first_item_str.startswith("data: "):
-                first_item_str = first_item_str[6:]
-            elif first_item_str.startswith("data:"):
-                first_item_str = first_item_str[5:]
-            first_item_str = json.loads(first_item_str)
-        if isinstance(first_item_str, dict) and 'error' in first_item_str:
-            print('\033[31m')
-            print(f"first_item_str: {first_item_str}")
-            print('\033[0m')
-            # 如果第一个 yield 的项是错误信息，抛出 HTTPException
-            raise HTTPException(status_code=status_code, detail=first_item_str)
-
-        # 如果不是错误，创建一个新的生成器，首先yield第一个项，然后yield剩余的项
-        async def new_generator():
-            yield first_item
-            async for item in generator:
-                yield item
-
-        return new_generator()
-    except StopAsyncIteration:
-        # 处理生成器为空的情况
-        return []
-
 async def process_request(request: RequestModel, provider: Dict):
     print("provider: ", provider['provider'])
     url = provider['base_url']
@@ -201,94 +140,10 @@ class ModelRequestHandler:
 
 model_handler = ModelRequestHandler()
 
-@app.middleware("http")
-async def log_requests(request: Request, call_next):
-    # 打印请求信息
-    logging.info(f"Request: {request.method} {request.url}")
-    # 打印请求体（如果有）
-    if request.method in ["POST", "PUT", "PATCH"]:
-        body = await request.body()
-        logging.info(f"Request Body: {body.decode('utf-8')}")
-
-    response = await call_next(request)
-    return response
-
-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
-
 @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)
 
-def get_all_models(token):
-    all_models = []
-    unique_models = set()
-
-    if token not in api_list:
-        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
-    api_index = api_list.index(token)
-    if config['api_keys'][api_index]['model']:
-        for model in config['api_keys'][api_index]['model']:
-            if "/" in model:
-                provider = model.split("/")[0]
-                model = model.split("/")[1]
-                if model == "*":
-                    for provider_item in config["providers"]:
-                        if provider_item['provider'] != provider:
-                            continue
-                        for model_item in provider_item['model'].keys():
-                            if model_item not in unique_models:
-                                unique_models.add(model_item)
-                                model_info = {
-                                    "id": model_item,
-                                    "object": "model",
-                                    "created": 1720524448858,
-                                    "owned_by": provider_item['provider']
-                                }
-                                all_models.append(model_info)
-                else:
-                    for provider_item in config["providers"]:
-                        if provider_item['provider'] != provider:
-                            continue
-                        for model_item in provider_item['model'].keys() :
-                            if model_item not in unique_models and model_item == model:
-                                unique_models.add(model_item)
-                                model_info = {
-                                    "id": model_item,
-                                    "object": "model",
-                                    "created": 1720524448858,
-                                    "owned_by": provider_item['provider']
-                                }
-                                all_models.append(model_info)
-                continue
-
-            if model not in unique_models:
-                unique_models.add(model)
-                model_info = {
-                    "id": model,
-                    "object": "model",
-                    "created": 1720524448858,
-                    "owned_by": model
-                }
-                all_models.append(model_info)
-    else:
-        for provider in config["providers"]:
-            for model in provider['model'].keys():
-                if model not in unique_models:
-                    unique_models.add(model)
-                    model_info = {
-                        "id": model,
-                        "object": "model",
-                        "created": 1720524448858,
-                        "owned_by": provider['provider']
-                    }
-                    all_models.append(model_info)
-
-    return all_models
-
 @app.post("/v1/models")
 async def list_models(token: str = Depends(verify_api_key)):
     models = get_all_models(token)

utils.py

@@ -0,0 +1,136 @@
+import yaml
+import json
+from fastapi import HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+# 读取YAML配置文件
+def load_config():
+    try:
+        with open('./api.yaml', 'r') as f:
+            conf = yaml.safe_load(f)
+            for index, provider in enumerate(conf['providers']):
+                model_dict = {}
+                for model in provider['model']:
+                    if type(model) == str:
+                        model_dict[model] = model
+                    if type(model) == dict:
+                        model_dict.update({value: key for key, value in model.items()})
+                provider['model'] = model_dict
+                conf['providers'][index] = provider
+            api_keys_db = conf['api_keys']
+            api_list = [item["api"] for item in api_keys_db]
+            # print(json.dumps(conf, indent=4, ensure_ascii=False))
+            return conf, api_keys_db, api_list
+    except FileNotFoundError:
+        print("配置文件 'api.yaml' 未找到。请确保文件存在于正确的位置。")
+        return [], [], []
+    except yaml.YAMLError:
+        print("配置文件 'api.yaml' 格式不正确。请检查 YAML 格式。")
+        return [], [], []
+
+config, api_keys_db, api_list = load_config()
+
+async def error_handling_wrapper(generator, status_code=200):
+    try:
+        first_item = await generator.__anext__()
+        first_item_str = first_item
+        if isinstance(first_item_str, (bytes, bytearray)):
+            first_item_str = first_item_str.decode("utf-8")
+        if isinstance(first_item_str, str):
+            if first_item_str.startswith("data: "):
+                first_item_str = first_item_str[6:]
+            elif first_item_str.startswith("data:"):
+                first_item_str = first_item_str[5:]
+            first_item_str = json.loads(first_item_str)
+        if isinstance(first_item_str, dict) and 'error' in first_item_str:
+            print('\033[31m')
+            print(f"first_item_str: {first_item_str}")
+            print('\033[0m')
+            # 如果第一个 yield 的项是错误信息，抛出 HTTPException
+            raise HTTPException(status_code=status_code, detail=first_item_str)
+
+        # 如果不是错误，创建一个新的生成器，首先yield第一个项，然后yield剩余的项
+        async def new_generator():
+            yield first_item
+            async for item in generator:
+                yield item
+
+        return new_generator()
+    except StopAsyncIteration:
+        # 处理生成器为空的情况
+        return []
+
+def get_all_models(token):
+    all_models = []
+    unique_models = set()
+
+    if token not in api_list:
+        raise HTTPException(status_code=403, detail="Invalid or missing API Key")
+    api_index = api_list.index(token)
+    if config['api_keys'][api_index]['model']:
+        for model in config['api_keys'][api_index]['model']:
+            if "/" in model:
+                provider = model.split("/")[0]
+                model = model.split("/")[1]
+                if model == "*":
+                    for provider_item in config["providers"]:
+                        if provider_item['provider'] != provider:
+                            continue
+                        for model_item in provider_item['model'].keys():
+                            if model_item not in unique_models:
+                                unique_models.add(model_item)
+                                model_info = {
+                                    "id": model_item,
+                                    "object": "model",
+                                    "created": 1720524448858,
+                                    "owned_by": provider_item['provider']
+                                }
+                                all_models.append(model_info)
+                else:
+                    for provider_item in config["providers"]:
+                        if provider_item['provider'] != provider:
+                            continue
+                        for model_item in provider_item['model'].keys() :
+                            if model_item not in unique_models and model_item == model:
+                                unique_models.add(model_item)
+                                model_info = {
+                                    "id": model_item,
+                                    "object": "model",
+                                    "created": 1720524448858,
+                                    "owned_by": provider_item['provider']
+                                }
+                                all_models.append(model_info)
+                continue
+
+            if model not in unique_models:
+                unique_models.add(model)
+                model_info = {
+                    "id": model,
+                    "object": "model",
+                    "created": 1720524448858,
+                    "owned_by": model
+                }
+                all_models.append(model_info)
+    else:
+        for provider in config["providers"]:
+            for model in provider['model'].keys():
+                if model not in unique_models:
+                    unique_models.add(model)
+                    model_info = {
+                        "id": model,
+                        "object": "model",
+                        "created": 1720524448858,
+                        "owned_by": provider['provider']
+                    }
+                    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