File size: 10,375 Bytes
d62232c 29e15d9 82cedb5 d62232c a1d3f78 810662a 93aef2a 810662a d62232c 3e22716 d62232c 077bdaa 728ade9 3e22716 728ade9 3e22716 465ab6f 3e22716 d62232c a1d3f78 3e22716 d62232c a1d3f78 3e22716 d62232c e6f7928 e6673ca d62232c 55a4247 e6673ca 55a4247 d62232c 29e15d9 133f622 e6f7928 eff878c 133f622 d62232c 1ad99a8 d62232c 55a4247 cf121cf e6673ca 55a4247 82cedb5 d62232c e7d701b 55a4247 51e9d12 55a4247 d62232c 3e22716 d62232c aa42078 d62232c aa42078 d62232c aa42078 3e22716 aa42078 d62232c 3e22716 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 |
import json
from fastapi import HTTPException
import httpx
from log_config import logger
def update_config(config_data):
for index, provider in enumerate(config_data['providers']):
model_dict = {}
for model in provider['model']:
if type(model) == str:
model_dict[model] = model
if type(model) == dict:
model_dict.update({new: old for old, new in model.items()})
provider['model'] = model_dict
config_data['providers'][index] = provider
api_keys_db = config_data['api_keys']
api_list = [item["api"] for item in api_keys_db]
# logger.info(json.dumps(config_data, indent=4, ensure_ascii=False))
return config_data, api_keys_db, api_list
# 读取YAML配置文件
async def load_config(app):
import yaml
try:
with open('./api.yaml', 'r') as f:
# 判断是否为空文件
conf = yaml.safe_load(f)
# conf = None
if conf:
config, api_keys_db, api_list = update_config(conf)
else:
# logger.error("配置文件 'api.yaml' 为空。请检查文件内容。")
config, api_keys_db, api_list = [], [], []
except FileNotFoundError:
logger.error("配置文件 'api.yaml' 未找到。请确保文件存在于正确的位置。")
config, api_keys_db, api_list = [], [], []
except yaml.YAMLError:
logger.error("配置文件 'api.yaml' 格式不正确。请检查 YAML 格式。")
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)):
return item.decode("utf-8")
elif isinstance(item, str):
return item
elif isinstance(item, dict):
return f"data: {json.dumps(item)}\n\n"
else:
return str(item)
import asyncio
async def error_handling_wrapper(generator, status_code=200):
async def new_generator():
try:
yield ensure_string(first_item)
async for item in generator:
yield ensure_string(item)
except (httpx.ReadError, asyncio.CancelledError) as e:
logger.error(f"Network error in new_generator: {e}")
raise HTTPException(status_code=503, detail=f"Stream interrupted: {str(e)}")
except Exception as e:
logger.exception(f"Error in new_generator: {e}")
raise HTTPException(status_code=status_code, detail=f"Stream error: {str(e)}")
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:]
if first_item_str.startswith("[DONE]"):
logger.error("error_handling_wrapper [DONE]!")
raise StopAsyncIteration
try:
first_item_str = json.loads(first_item_str)
except json.JSONDecodeError:
logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
raise StopAsyncIteration
if isinstance(first_item_str, dict) and 'error' in first_item_str:
raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
wrapped_generator = new_generator()
try:
async for item in wrapped_generator:
yield item
except HTTPException as http_exc:
raise HTTPException(status_code=status_code, detail=f"Wrapper error: {str(http_exc)}")
except (httpx.ReadError, asyncio.CancelledError) as e:
logger.error(f"Network error during streaming: {e}")
raise HTTPException(status_code=503, detail=f"Stream interrupted: {str(e)}")
except Exception as e:
logger.exception(f"Unexpected error in error_handling_wrapper: {e}")
raise HTTPException(status_code=status_code, detail=f"Unexpected error: {str(e)}")
except StopAsyncIteration:
raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
except HTTPException as http_exc:
raise HTTPException(status_code=status_code, detail=f"Wrapper error: {str(http_exc)}")
except Exception as e:
logger.exception(f"Error in error_handling_wrapper: {e}")
raise HTTPException(status_code=status_code, detail=f"Wrapper error: {str(e)}")
# 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:]
# if first_item_str.startswith("[DONE]"):
# logger.error("error_handling_wrapper [DONE]!")
# raise StopAsyncIteration
# try:
# first_item_str = json.loads(first_item_str)
# except json.JSONDecodeError:
# logger.error("error_handling_wrapper JSONDecodeError!" + repr(first_item_str))
# raise StopAsyncIteration
# if isinstance(first_item_str, dict) and 'error' in first_item_str:
# # 如果第一个 yield 的项是错误信息，抛出 HTTPException
# raise HTTPException(status_code=status_code, detail=f"{first_item_str}"[:300])
# # 如果不是错误，创建一个新的生成器，首先yield第一个项，然后yield剩余的项
# async def new_generator():
# yield ensure_string(first_item)
# async for item in generator:
# yield ensure_string(item)
# return new_generator()
# except StopAsyncIteration:
# raise HTTPException(status_code=status_code, detail="data: {'error': 'No data returned'}")
def post_all_models(token, config, api_list):
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": "uni-api"
# "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": "uni-api"
}
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)
return all_models
def get_all_models(config):
all_models = []
unique_models = set()
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": "uni-api"
}
all_models.append(model_info)
return all_models |