import asyncio
import json
import random # 添加 random 模块导入
import re
from datetime import datetime
from typing import Tuple, Union, Dict, Any, Set # 引入 Set
import aiohttp
from aiohttp.client import ClientResponse
# 相对路径导入，根据你的项目结构调整
# 例如，如果 utils_model.py 在 src/utils/ 下，而 logger 在 src/common/ 下
# from ..common.logger import get_module_logger
# from ..common.database import db
# from ..config.config import global_config
# 假设它们在期望的路径
from src.common.logger import get_module_logger
from ...common.database import db
from ...config.config import global_config
import base64
from PIL import Image
import io
import os
from rich.traceback import install
install(extra_lines=3)
# 尝试加载 .env 文件中的环境变量 (如果项目结构需要)
# load_dotenv() # 如果你的 .env 文件不在标准位置，可能需要指定路径 load_dotenv(dotenv_path='path/to/.env')
logger = get_module_logger("model_utils")
class PayLoadTooLargeError(Exception):
"""自定义异常类，用于处理请求体过大错误"""
# (代码不变)
def __init__(self, message: str):
super().__init__(message)
self.message = message
def __str__(self):
return "请求体过大，请尝试压缩图片或减少输入内容。"
class RequestAbortException(Exception):
"""自定义异常类，用于处理请求中断异常"""
# (代码不变)
def __init__(self, message: str, response: ClientResponse):
super().__init__(message)
self.message = message
self.response = response
def __str__(self):
return self.message
class PermissionDeniedException(Exception):
"""自定义异常类，用于处理访问拒绝的异常"""
# (代码不变)
def __init__(self, message: str, key_identifier: str = None): # 添加 key 标识符
super().__init__(message)
self.message = message
self.key_identifier = key_identifier # 存储导致 403 的 key
def __str__(self):
return self.message
# 新增：用于内部标记需要切换 Key 的异常
class _SwitchKeyException(Exception):
"""内部异常，用于标记需要切换Key并且跳过标准等待时间."""
# (代码不变)
pass
# 常见Error Code Mapping
error_code_mapping = {
# (代码不变)
400: "参数不正确",
401: "API key 错误，认证失败，请检查/config/bot_config.toml和.env中的配置是否正确哦~", # 401 也可能是 Key 无效
402: "账号余额不足",
403: "需要实名,或余额不足,或Key无权限", # 扩展 403 的含义
404: "Not Found",
429: "请求过于频繁，请稍后再试",
500: "服务器内部故障",
503: "服务器负载过高",
}
async def _safely_record(request_content: Dict[str, Any], payload: Dict[str, Any]):
"""安全地记录请求内容，隐藏敏感信息"""
# (代码不变)
image_base64: str = request_content.get("image_base64")
image_format: str = request_content.get("image_format")
is_gemini_payload = payload and isinstance(payload, dict) and "contents" in payload
safe_payload = json.loads(json.dumps(payload)) if payload else {}
if image_base64 and safe_payload and isinstance(safe_payload, dict):
if "messages" in safe_payload and len(safe_payload["messages"]) > 0:
if isinstance(safe_payload["messages"][0], dict) and "content" in safe_payload["messages"][0]:
content = safe_payload["messages"][0]["content"]
if (
isinstance(content, list)
and len(content) > 1
and isinstance(content[1], dict)
and "image_url" in content[1]
):
safe_payload["messages"][0]["content"][1]["image_url"]["url"] = (
f"data:image/{image_format.lower() if image_format else 'jpeg'};base64,"
f"{image_base64[:10]}...{image_base64[-10:]}"
)
elif is_gemini_payload and "contents" in safe_payload and len(safe_payload["contents"]) > 0:
if isinstance(safe_payload["contents"][0], dict) and "parts" in safe_payload["contents"][0]:
parts = safe_payload["contents"][0]["parts"]
for i, part in enumerate(parts):
if isinstance(part, dict) and "inlineData" in part:
safe_payload["contents"][0]["parts"][i]["inlineData"]["data"] = (
f"{image_base64[:10]}...{image_base64[-10:]}"
)
break
return safe_payload
class LLMRequest:
# (代码不变)
MODELS_NEEDING_TRANSFORMATION = [
"o1",
"o1-2024-12-17",
"o1-mini",
"o1-mini-2024-09-12",
"o1-preview",
"o1-preview-2024-09-12",
"o1-pro",
"o1-pro-2025-03-19",
"o3",
"o3-2025-04-16",
"o3-mini",
"o3-mini-2025-01-31o4-mini",
"o4-mini-2025-04-16",
]
_abandoned_keys_runtime: Set[str] = set()
def __init__(self, model: dict, **kwargs):
"""初始化 LLMRequest 实例"""
# (代码不变)
self.model_key_name = model["key"]
self.model_name: str = model["name"]
self.params = kwargs
self.stream = model.get("stream", False)
self.pri_in = model.get("pri_in", 0)
self.pri_out = model.get("pri_out", 0)
self.request_type = model.get("request_type", "default")
try:
raw_api_key_config = os.environ[self.model_key_name]
self.base_url = os.environ[model["base_url"]]
self.is_gemini = "googleapis.com" in self.base_url.lower()
if self.is_gemini:
logger.debug(f"模型 {self.model_name}: 检测到为 Gemini API (Base URL: {self.base_url})")
if self.stream:
logger.warning(f"模型 {self.model_name}: Gemini 流式输出处理与 OpenAI 不同，暂时强制禁用流式。")
self.stream = False
# 解析和过滤 API Keys (代码不变)
parsed_keys = []
is_list_config = False
try:
loaded_keys = json.loads(raw_api_key_config)
if isinstance(loaded_keys, list):
parsed_keys = [str(key) for key in loaded_keys if key]
is_list_config = True
elif isinstance(loaded_keys, str) and loaded_keys:
parsed_keys = [loaded_keys]
else:
raise ValueError(f"Parsed API key for {self.model_key_name} is not a valid list or string.")
except (json.JSONDecodeError, TypeError) as e:
if isinstance(raw_api_key_config, list):
parsed_keys = [str(key) for key in raw_api_key_config if key]
is_list_config = True
elif isinstance(raw_api_key_config, str) and raw_api_key_config:
parsed_keys = [raw_api_key_config]
else:
raise ValueError(
f"Invalid or empty API key config for {self.model_key_name}: {raw_api_key_config}"
) from e
if not parsed_keys:
raise ValueError(f"No valid API keys found for {self.model_key_name}.")
abandoned_key_name = f"abandon_{self.model_key_name}"
abandoned_keys_set = set()
raw_abandoned_keys = os.environ.get(abandoned_key_name)
if raw_abandoned_keys:
try:
loaded_abandoned = json.loads(raw_abandoned_keys)
if isinstance(loaded_abandoned, list):
abandoned_keys_set.update(str(key) for key in loaded_abandoned if key)
elif isinstance(loaded_abandoned, str) and loaded_abandoned:
abandoned_keys_set.add(loaded_abandoned)
logger.info(
f"模型 {model['name']}: 加载了 {len(abandoned_keys_set)} 个来自配置 '{abandoned_key_name}' 的废弃 Keys。"
)
except (json.JSONDecodeError, TypeError):
if isinstance(raw_abandoned_keys, list):
abandoned_keys_set.update(str(key) for key in raw_abandoned_keys if key)
logger.info(
f"模型 {model['name']}: 加载了 {len(abandoned_keys_set)} 个来自配置 '{abandoned_key_name}' (直接列表) 的废弃 Keys。"
)
elif isinstance(raw_abandoned_keys, str) and raw_abandoned_keys:
abandoned_keys_set.add(raw_abandoned_keys)
logger.info(
f"模型 {model['name']}: 加载了 1 个来自配置 '{abandoned_key_name}' (字符串) 的废弃 Key。"
)
else:
logger.warning(f"无法解析环境变量 '{abandoned_key_name}' 的内容: {raw_abandoned_keys}")
all_abandoned_keys = abandoned_keys_set.union(LLMRequest._abandoned_keys_runtime)
active_keys = [key for key in parsed_keys if key not in all_abandoned_keys]
if not active_keys:
logger.error(f"模型 {model['name']}: 所有为 '{self.model_key_name}' 配置的 Keys 都已被废弃或无效。")
raise ValueError(
f"No active API keys available for {self.model_key_name} after filtering abandoned keys."
)
if is_list_config and len(active_keys) > 1:
self._api_key_config = active_keys
logger.info(
f"模型 {model['name']}: 初始化完成，可用 Keys: {len(self._api_key_config)} (已排除 {len(all_abandoned_keys)} 个废弃 Keys)。"
)
elif active_keys:
self._api_key_config = active_keys[0]
logger.info(
f"模型 {model['name']}: 初始化完成，使用单个活动 Key (已排除 {len(all_abandoned_keys)} 个废弃 Keys)。"
)
else:
raise ValueError(f"Unexpected state: No active keys for {self.model_key_name}.")
# 加载代理配置 (代码不变)
self.proxy_url = None
self.proxy_models_set = set()
proxy_host = os.environ.get("PROXY_HOST")
proxy_port = os.environ.get("PROXY_PORT")
proxy_models_str = os.environ.get("PROXY_MODELS", "")
if proxy_host and proxy_port:
try:
int(proxy_port)
self.proxy_url = f"http://{proxy_host}:{proxy_port}"
logger.debug(f"代理已配置: {self.proxy_url}")
if proxy_models_str:
try:
cleaned_str = proxy_models_str.strip("'\"")
self.proxy_models_set = {
model_name.strip() for model_name in cleaned_str.split(",") if model_name.strip()
}
logger.debug(f"以下模型将使用代理: {self.proxy_models_set}")
except Exception as e:
logger.error(
f"解析 PROXY_MODELS ('{proxy_models_str}') 出错: {e}. 代理将不会对特定模型生效。"
)
self.proxy_models_set = set()
except ValueError:
logger.error(f"无效的代理端口号: {proxy_port}。代理将不被启用。")
self.proxy_url = None
self.proxy_models_set = set()
except Exception as e:
logger.error(f"加载代理配置时发生错误: {e}")
self.proxy_url = None
self.proxy_models_set = set()
else:
logger.info("未配置代理服务器 (PROXY_HOST 或 PROXY_PORT 未设置)。")
except KeyError as e:
# (代码不变)
missing_key = str(e).strip("'")
if missing_key == self.model_key_name:
logger.error(f"配置错误：找不到 API Key 环境变量 '{self.model_key_name}'")
raise ValueError(f"配置错误：找不到 API Key 环境变量 '{self.model_key_name}'") from e
elif missing_key == model["base_url"]:
logger.error(f"配置错误：找不到 Base URL 环境变量 '{model['base_url']}'")
raise ValueError(f"配置错误：找不到 Base URL 环境变量 '{model['base_url']}'") from e
else:
logger.error(f"配置错误：找不到环境变量 - {str(e)}")
raise ValueError(f"配置错误：找不到环境变量 - {str(e)}") from e
except AttributeError as e:
# (代码不变)
logger.error(f"原始 model dict 信息：{model}")
logger.error(f"配置错误：找不到对应的配置项 - {str(e)}")
raise ValueError(f"配置错误：找不到对应的配置项 - {str(e)}") from e
except ValueError as e:
# (代码不变)
logger.error(f"API Key 或配置初始化错误 for {self.model_key_name}: {str(e)}")
raise e
self._init_database()
@staticmethod
def _init_database():
"""初始化数据库集合"""
# (代码不变)
try:
db.llm_usage.create_index([("timestamp", 1)])
db.llm_usage.create_index([("model_name", 1)])
db.llm_usage.create_index([("user_id", 1)])
db.llm_usage.create_index([("request_type", 1)])
except Exception as e:
logger.error(f"创建数据库索引失败: {str(e)}")
def _record_usage(
self,
prompt_tokens: int,
completion_tokens: int,
total_tokens: int,
user_id: str = "system",
request_type: str = None,
endpoint: str = "/chat/completions",
):
"""记录模型使用情况到数据库
Args:
prompt_tokens: 输入token数
completion_tokens: 输出token数
total_tokens: 总token数
user_id: 用户ID，默认为system
request_type: 请求类型
endpoint: API端点
"""
# 如果 request_type 为 None，则使用实例变量中的值
if request_type is None:
request_type = self.request_type
actual_endpoint = endpoint
if self.is_gemini:
if endpoint == "/embeddings":
actual_endpoint = ":embedContent"
else:
actual_endpoint = ":generateContent"
try:
usage_data = {
"model_name": self.model_name,
"user_id": user_id,
"request_type": request_type or self.request_type,
"endpoint": actual_endpoint,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
"cost": self._calculate_cost(prompt_tokens, completion_tokens),
"status": "success",
"timestamp": datetime.now(),
}
db.llm_usage.insert_one(usage_data)
logger.trace(
f"Token使用情况 - 模型: {self.model_name}, "
f"用户: {user_id}, 类型: {request_type or self.request_type}, "
f"提示词: {prompt_tokens}, 完成: {completion_tokens}, "
f"总计: {total_tokens}"
)
except Exception as e:
logger.error(f"记录token使用情况失败: {str(e)}")
def _calculate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
"""计算API调用成本"""
# (代码不变)
input_cost = (prompt_tokens / 1000000) * self.pri_in
output_cost = (completion_tokens / 1000000) * self.pri_out
return round(input_cost + output_cost, 6)
async def _prepare_request(
self,
endpoint: str,
prompt: str = None,
image_base64: str = None,
image_format: str = None,
payload: dict = None,
retry_policy: dict = None,
**kwargs: Any,
) -> Dict[str, Any]:
"""配置请求参数，合并实例参数和调用时参数"""
default_retry = {
"max_retries": global_config.api_polling_max_retries,
"base_wait": 10,
"retry_codes": [429, 413, 500, 503],
"abort_codes": [400, 401, 402, 403],
}
policy = {**default_retry, **(retry_policy or {})}
_actual_endpoint = endpoint
if self.is_gemini:
action = endpoint.lstrip("/")
api_url = f"{self.base_url.rstrip('/')}/{self.model_name}{action}"
stream_mode = False
else:
api_url = f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}"
stream_mode = self.stream
call_params = {k: v for k, v in kwargs.items() if k != "request_type"}
merged_params = {**self.params, **call_params}
if payload is None:
payload = await self._build_payload(prompt, image_base64, image_format, merged_params)
else:
logger.debug("使用外部提供的 payload，忽略单次调用参数合并。")
if not self.is_gemini and stream_mode:
payload["stream"] = merged_params.get("stream", stream_mode)
return {
"policy": policy,
"payload": payload,
"api_url": api_url,
"stream_mode": payload.get("stream", False),
"image_base64": image_base64,
"image_format": image_format,
"prompt": prompt,
}
async def _execute_request(
self,
endpoint: str,
prompt: str = None,
image_base64: str = None,
image_format: str = None,
payload: dict = None,
retry_policy: dict = None,
response_handler: callable = None,
user_id: str = "system",
request_type: str = None,
**kwargs: Any,
):
"""统一请求执行入口, 支持列表 key 切换、代理和单次调用参数覆盖"""
final_request_type = request_type or kwargs.get("request_type") or self.request_type
api_kwargs = {k: v for k, v in kwargs.items() if k != "request_type"}
request_content = await self._prepare_request(
endpoint, prompt, image_base64, image_format, payload, retry_policy, **api_kwargs
)
policy = request_content["policy"]
api_url = request_content["api_url"]
actual_payload = request_content["payload"]
stream_mode = request_content["stream_mode"]
use_proxy = False
current_proxy_url = None
if self.proxy_url and self.model_name in self.proxy_models_set:
use_proxy = True
current_proxy_url = self.proxy_url
logger.debug(f"模型 {self.model_name}: 将通过代理 {current_proxy_url} 发送请求。")
elif self.proxy_url:
logger.debug(f"模型 {self.model_name}: 配置了代理，但此模型不在 PROXY_MODELS 列表中，将不使用代理。")
else:
logger.debug(f"模型 {self.model_name}: 未配置或不为此模型使用代理。")
current_key = None
keys_failed_429 = set()
keys_abandoned_runtime = set()
key_switch_limit_429 = global_config.api_polling_max_retries
key_switch_limit_403 = global_config.api_polling_max_retries
available_keys_pool = []
is_key_list = isinstance(self._api_key_config, list)
if is_key_list:
available_keys_pool = list(self._api_key_config)
if not available_keys_pool:
logger.error(f"模型 {self.model_name}: 初始化后无可用活动 Keys。")
raise ValueError(f"模型 {self.model_name}: 无可用活动 Keys。")
random.shuffle(available_keys_pool)
key_switch_limit_429 = min(key_switch_limit_429, len(available_keys_pool))
key_switch_limit_403 = min(key_switch_limit_403, len(available_keys_pool))
logger.info(
f"模型 {self.model_name}: Key 列表模式，启用 429/403 自动切换（429上限: {key_switch_limit_429}, 403上限: {key_switch_limit_403}）。"
)
elif isinstance(self._api_key_config, str):
available_keys_pool = [self._api_key_config]
key_switch_limit_429 = 1
key_switch_limit_403 = 1
else:
logger.error(f"模型 {self.model_name}: 无效的 API Key 配置类型在执行时遇到: {type(self._api_key_config)}")
raise TypeError(f"模型 {self.model_name}: 无效的 API Key 配置类型")
last_exception = None
for attempt in range(policy["max_retries"]):
if available_keys_pool:
current_key = available_keys_pool.pop(0)
elif current_key:
logger.debug(
f"模型 {self.model_name}: 无新 Key 可用或为单 Key 模式，将使用 Key ...{current_key[-4:]} 进行重试 (第 {attempt + 1} 次尝试)"
)
else:
if (
not self._api_key_config
or all(
k in LLMRequest._abandoned_keys_runtime
for k in self._api_key_config
if isinstance(self._api_key_config, list)
)
or (
isinstance(self._api_key_config, str)
and self._api_key_config in LLMRequest._abandoned_keys_runtime
)
):
final_error_msg = f"模型 {self.model_name}: 所有可用 API Keys 均因 403 错误被禁用。"
logger.critical(final_error_msg)
raise PermissionDeniedException(final_error_msg)
else:
raise RuntimeError(f"模型 {self.model_name}: 无法选择 API key (第 {attempt + 1} 次尝试)")
logger.debug(f"模型 {self.model_name}: 尝试使用 Key: ...{current_key[-4:]} (总第 {attempt + 1} 次尝试)")
try:
headers = await self._build_headers(current_key)
if not self.is_gemini and stream_mode:
headers["Accept"] = "text/event-stream"
async with aiohttp.ClientSession() as session:
post_kwargs = {"headers": headers, "json": actual_payload, "timeout": 60}
if use_proxy:
post_kwargs["proxy"] = current_proxy_url
async with session.post(api_url, **post_kwargs) as response:
if response.status == 429 and is_key_list:
logger.warning(f"模型 {self.model_name}: Key ...{current_key[-4:]} 遇到 429 错误。")
response_text = await response.text()
logger.debug(
f"模型 {self.model_name}: Key ...{current_key[-4:]} response:\n{json.dumps(json.loads(response_text), indent=2, ensure_ascii=False)}\napi_url:\n{api_url}\nheader:\n{headers}\npayload:\n{actual_payload}"
)
if current_key not in keys_failed_429:
keys_failed_429.add(current_key)
logger.info(
f" (因 429 已失败 {len(keys_failed_429)}/{key_switch_limit_429} 个不同 Key)"
)
if available_keys_pool and len(keys_failed_429) < key_switch_limit_429:
logger.info(" 尝试因 429 切换到下一个可用 Key...")
raise _SwitchKeyException()
else:
logger.warning(" 无更多 Key 可因 429 切换或已达上限。")
else:
logger.warning(f" Key ...{current_key[-4:]} 再次遇到 429，按标准重试流程。")
elif response.status == 403 and is_key_list:
logger.error(
f"模型 {self.model_name}: Key ...{current_key[-4:]} 遇到 403 (权限拒绝) 错误。"
)
if current_key not in keys_abandoned_runtime:
keys_abandoned_runtime.add(current_key)
LLMRequest._abandoned_keys_runtime.add(current_key)
logger.critical(
f" !! Key ...{current_key[-4:]} 已添加到运行时废弃列表。请考虑将其移至配置中的 'abandon_{self.model_key_name}' !!"
)
if current_key in available_keys_pool:
available_keys_pool.remove(current_key)
if available_keys_pool and len(keys_abandoned_runtime) < key_switch_limit_403:
logger.info(" 尝试因 403 切换到下一个可用 Key...")
raise _SwitchKeyException()
else:
logger.error(" 无更多 Key 可因 403 切换或已达上限。将中止请求。")
await response.read()
raise PermissionDeniedException(
f"Key ...{current_key[-4:]} 权限被拒，且无其他可用 Key 切换。",
key_identifier=current_key,
)
else:
logger.error(f" Key ...{current_key[-4:]} 再次遇到 403，这不应发生。中止请求。")
await response.read()
raise PermissionDeniedException(
f"Key ...{current_key[-4:]} 重复遇到 403。", key_identifier=current_key
)
elif response.status in policy["retry_codes"] or response.status in policy["abort_codes"]:
await self._handle_error_response(response, attempt, policy, current_key)
if response.status in policy["retry_codes"] and attempt < policy["max_retries"] - 1:
if response.status not in [429, 403]:
wait_time = policy["base_wait"] * (2**attempt)
logger.warning(
f"模型 {self.model_name}: 遇到可重试错误 {response.status}, 等待 {wait_time} 秒后重试..."
)
await asyncio.sleep(wait_time)
last_exception = RuntimeError(f"重试错误 {response.status}")
continue
if response.status in policy["abort_codes"] or (
response.status in policy["retry_codes"] and attempt >= policy["max_retries"] - 1
):
if attempt >= policy["max_retries"] - 1 and response.status in policy["retry_codes"]:
logger.error(
f"模型 {self.model_name}: 达到最大重试次数，最后一次尝试仍为可重试错误 {response.status}。"
)
# await self._handle_error_response(response, attempt, policy, current_key)
# await response.read()
# final_error_msg = f"请求中止或达到最大重试次数，最终状态码: {response.status}"
# logger.error(final_error_msg)
# raise RequestAbortException(final_error_msg, response)
response.raise_for_status()
result = {}
if not self.is_gemini and stream_mode:
result = await self._handle_stream_output(response)
else:
result = await response.json()
return (
response_handler(result)
if response_handler
else self._default_response_handler(result, user_id, final_request_type, endpoint)
)
except _SwitchKeyException:
last_exception = _SwitchKeyException()
logger.debug("捕获到 _SwitchKeyException，立即进行下一次尝试。")
continue
except PermissionDeniedException as e:
logger.error(f"模型 {self.model_name}: 因权限拒绝 (403) 中止请求: {e}")
if is_key_list and not available_keys_pool and e.key_identifier:
logger.critical(f" 中止原因是 Key ...{e.key_identifier[-4:]} 触发 403 后已无其他 Key 可用。")
raise e
except aiohttp.ClientProxyConnectionError as e:
logger.error(f"代理连接错误: {e} (代理地址: {current_proxy_url})")
last_exception = e
if attempt >= policy["max_retries"] - 1:
raise RuntimeError(f"代理连接失败达到最大重试次数: {e}") from e
wait_time = policy["base_wait"] * (2**attempt)
logger.warning(f"模型 {self.model_name}: 代理连接错误，等待 {wait_time} 秒后重试...")
await asyncio.sleep(wait_time)
continue
except aiohttp.ClientConnectorError as e:
logger.error(f"网络连接错误: {e} (URL: {api_url}, 代理: {current_proxy_url})")
last_exception = e
if attempt >= policy["max_retries"] - 1:
raise RuntimeError(f"网络连接失败达到最大重试次数: {e}") from e
wait_time = policy["base_wait"] * (2**attempt)
logger.warning(f"模型 {self.model_name}: 网络连接错误，等待 {wait_time} 秒后重试...")
await asyncio.sleep(wait_time)
continue
except (PayLoadTooLargeError, RequestAbortException) as e:
# (代码不变)
logger.error(f"模型 {self.model_name}: 请求处理中遇到关键错误，将中止: {e}")
raise e
except Exception as e:
# (代码不变)
last_exception = e
logger.warning(
f"模型 {self.model_name}: 第 {attempt + 1} 次尝试中发生非 HTTP 错误: {str(e.__class__.__name__)} - {str(e)}"
)
if attempt >= policy["max_retries"] - 1:
logger.error(
f"模型 {self.model_name}: 达到最大重试次数 ({policy['max_retries']})，因非 HTTP 错误失败。"
)
else:
try:
temp_request_content = {
"policy": policy,
"payload": actual_payload,
"api_url": api_url,
"stream_mode": stream_mode,
"image_base64": image_base64,
"image_format": image_format,
"prompt": prompt,
}
handled_payload, count_delta = await self._handle_exception(
e, attempt, temp_request_content, merged_params=api_kwargs
)
if handled_payload:
actual_payload = handled_payload
logger.info(f"模型 {self.model_name}: 异常处理更新了 payload，将使用当前 Key 重试。")
wait_time = policy["base_wait"] * (2**attempt)
logger.warning(f"模型 {self.model_name}: 等待 {wait_time} 秒后重试...")
await asyncio.sleep(wait_time)
continue
except (RequestAbortException, PermissionDeniedException) as abort_exception:
logger.error(f"模型 {self.model_name}: 异常处理判断需要中止请求: {abort_exception}")
raise abort_exception
except RuntimeError as rt_error:
logger.error(f"模型 {self.model_name}: 异常处理遇到运行时错误: {rt_error}")
raise rt_error
# --- 循环结束 ---
logger.error(f"模型 {self.model_name}: 所有重试尝试 ({policy['max_retries']} 次) 均失败。")
if last_exception:
if isinstance(last_exception, PermissionDeniedException):
logger.error(f"最后遇到的错误是权限拒绝: {str(last_exception)}")
raise last_exception
logger.error(f"最后遇到的错误: {str(last_exception.__class__.__name__)} - {str(last_exception)}")
raise RuntimeError(
f"模型 {self.model_name} 达到最大重试次数，API 请求失败。最后错误: {str(last_exception)}"
) from last_exception
else:
if not available_keys_pool and keys_abandoned_runtime:
final_error_msg = f"模型 {self.model_name}: 所有可用 API Keys 均因 403 错误被禁用。"
logger.critical(final_error_msg)
raise PermissionDeniedException(final_error_msg)
else:
raise RuntimeError(f"模型 {self.model_name} 达到最大重试次数，API 请求失败，原因未知。")
async def _handle_stream_output(self, response: ClientResponse) -> Dict[str, Any]:
"""处理 OpenAI 兼容的流式输出"""
# (代码不变)
flag_delta_content_finished = False
accumulated_content = ""
usage = None
reasoning_content = ""
content = ""
tool_calls = None
async for line_bytes in response.content:
try:
line = line_bytes.decode("utf-8").strip()
if not line:
continue
if line.startswith("data:"):
data_str = line[5:].strip()
if data_str == "[DONE]":
break
try:
chunk = json.loads(data_str)
if flag_delta_content_finished:
chunk_usage = chunk.get("usage", None)
if chunk_usage:
usage = chunk_usage
else:
delta = chunk["choices"][0]["delta"]
delta_content = delta.get("content")
if delta_content is None:
delta_content = ""
accumulated_content += delta_content
if "tool_calls" in delta:
if tool_calls is None:
tool_calls = []
for tc in delta["tool_calls"]:
new_tc = dict(tc)
if "function" in new_tc and "arguments" not in new_tc["function"]:
new_tc["function"]["arguments"] = ""
tool_calls.append(new_tc)
else:
for i, tc_delta in enumerate(delta["tool_calls"]):
if (
i < len(tool_calls)
and "function" in tc_delta
and "arguments" in tc_delta["function"]
):
if "arguments" in tool_calls[i]["function"]:
tool_calls[i]["function"]["arguments"] += tc_delta["function"][
"arguments"
]
else:
tool_calls[i]["function"]["arguments"] = tc_delta["function"][
"arguments"
]
elif i >= len(tool_calls):
new_tc = dict(tc_delta)
if "function" in new_tc and "arguments" not in new_tc["function"]:
new_tc["function"]["arguments"] = ""
tool_calls.append(new_tc)
finish_reason = chunk["choices"][0].get("finish_reason")
if delta.get("reasoning_content", None):
reasoning_content += delta["reasoning_content"]
if finish_reason == "stop" or finish_reason == "tool_calls":
chunk_usage = chunk.get("usage", None)
if chunk_usage:
usage = chunk_usage
break
flag_delta_content_finished = True
except json.JSONDecodeError as e:
logger.error(f"模型 {self.model_name} 解析流式 JSON 错误: {e} - data: '{data_str}'")
except Exception as e:
logger.exception(f"模型 {self.model_name} 解析流式输出块错误: {str(e)}")
except UnicodeDecodeError as e:
logger.warning(f"模型 {self.model_name} 流式输出解码错误: {e} - bytes: {line_bytes[:50]}...")
except Exception as e:
if isinstance(e, GeneratorExit):
log_content = f"模型 {self.model_name} 流式输出被中断，正在清理资源..."
else:
log_content = f"模型 {self.model_name} 处理流式输出时发生错误: {str(e)}"
logger.warning(log_content)
try:
await response.release()
except Exception as cleanup_error:
logger.error(f"清理资源时发生错误: {cleanup_error}")
content = accumulated_content
break
if not content and accumulated_content:
content = accumulated_content
think_match = re.search(r"(.*?)", content, re.DOTALL)
if think_match:
reasoning_content = think_match.group(1).strip()
content = re.sub(r".*?", "", content, flags=re.DOTALL).strip()
message = {
"content": content,
"reasoning_content": reasoning_content,
}
if tool_calls:
message["tool_calls"] = tool_calls
result = {
"choices": [{"message": message}],
"usage": usage,
}
return result
async def _handle_error_response(
self, response: ClientResponse, retry_count: int, policy: Dict[str, Any], current_key: str = None
) -> None:
"""处理 HTTP 错误响应 (区分 403 和其他错误)"""
# (代码不变)
status = response.status
try:
error_text = await response.text()
except Exception as e:
error_text = f"(无法读取响应体: {e})"
if status == 403:
logger.error(
f"模型 {self.model_name}: 遇到 403 (权限拒绝) 错误。Key: ...{current_key[-4:] if current_key else 'N/A'}. "
f"响应: {error_text[:200]}"
)
raise PermissionDeniedException(f"模型禁止访问 ({status})", key_identifier=current_key)
elif status in policy["retry_codes"] and status != 429:
if status == 413:
logger.warning(
f"模型 {self.model_name}: 错误码 413 (Payload Too Large)。Key: ...{current_key[-4:] if current_key else 'N/A'}. 尝试压缩..."
)
raise PayLoadTooLargeError("请求体过大")
elif status in [500, 503]:
logger.error(
f"模型 {self.model_name}: 服务器内部错误或过载 ({status})。Key: ...{current_key[-4:] if current_key else 'N/A'}. "
f"响应: {error_text[:200]}"
)
return
else:
logger.warning(
f"模型 {self.model_name}: 遇到可重试错误码: {status}. Key: ...{current_key[-4:] if current_key else 'N/A'}"
)
return
elif status in policy["abort_codes"]:
logger.error(
f"模型 {self.model_name}: 遇到需要中止的错误码: {status} - {error_code_mapping.get(status, '未知错误')}. "
f"Key: ...{current_key[-4:] if current_key else 'N/A'}. 响应: {error_text[:200]}"
)
raise RequestAbortException(f"请求出现错误 {status}，中止处理", response)
else:
logger.error(
f"模型 {self.model_name}: 遇到未明确处理的错误码: {status}. Key: ...{current_key[-4:] if current_key else 'N/A'}. 响应: {error_text[:200]}"
)
try:
response.raise_for_status()
raise RequestAbortException(f"未处理的错误状态码 {status}", response)
except aiohttp.ClientResponseError as e:
raise RequestAbortException(f"未处理的错误状态码 {status}: {e.message}", response) from e
async def _handle_exception(
self, exception, retry_count: int, request_content: Dict[str, Any], merged_params: Dict[str, Any] = None
) -> Union[Tuple[Dict[str, Any], int], Tuple[None, int]]:
"""处理非 HTTP 错误，支持使用合并后的参数重建 payload"""
policy = request_content["policy"]
payload = request_content["payload"]
_wait_time = policy["base_wait"] * (2**retry_count)
keep_request = False
if retry_count < policy["max_retries"] - 1:
keep_request = True
params_for_rebuild = merged_params if merged_params is not None else payload
if isinstance(exception, PayLoadTooLargeError):
if keep_request:
logger.warning("请求体过大 (PayLoadTooLargeError)，尝试压缩图片...")
image_base64 = request_content.get("image_base64")
if image_base64:
compressed_image_base64 = compress_base64_image_by_scale(image_base64)
if compressed_image_base64 != image_base64:
new_payload = await self._build_payload(
request_content["prompt"],
compressed_image_base64,
request_content["image_format"],
params_for_rebuild,
)
logger.info("图片压缩成功，将使用压缩后的图片重试。")
return new_payload, 0
else:
logger.warning("图片压缩未改变大小或失败。")
else:
logger.warning("请求体过大但请求中不包含图片，无法压缩。")
return None, 0
else:
logger.error("达到最大重试次数，请求体仍然过大。")
raise RuntimeError("请求体过大，压缩或重试后仍然失败。") from exception
elif isinstance(exception, (aiohttp.ClientError, asyncio.TimeoutError)):
if keep_request:
logger.error(f"模型 {self.model_name} 网络错误: {str(exception)}")
return None, 0
else:
logger.critical(f"模型 {self.model_name} 网络错误达到最大重试次数: {str(exception)}")
raise RuntimeError(f"网络请求失败: {str(exception)}") from exception
elif isinstance(exception, aiohttp.ClientResponseError):
if keep_request:
logger.error(
f"模型 {self.model_name} HTTP响应错误 (未被策略覆盖): 状态码: {exception.status}, 错误: {exception.message}"
)
try:
error_text = await exception.response.text() if hasattr(exception, "response") else str(exception)
logger.error(f"服务器错误响应详情: {error_text[:500]}")
except Exception as parse_err:
logger.warning(f"无法解析服务器错误响应内容: {str(parse_err)}")
return None, 0
else:
logger.critical(
f"模型 {self.model_name} HTTP响应错误达到最大重试次数: 状态码: {exception.status}, 错误: {exception.message}"
)
current_key_placeholder = request_content.get("current_key", "******")
handled_payload = await _safely_record(request_content, payload)
logger.critical(
f"请求头: {await self._build_headers(api_key=current_key_placeholder, no_key=True)} 请求体: {handled_payload}"
)
raise RuntimeError(
f"模型 {self.model_name} API请求失败: 状态码 {exception.status}, {exception.message}"
) from exception
else:
if keep_request:
logger.error(
f"模型 {self.model_name} 遇到未知错误: {str(exception.__class__.__name__)} - {str(exception)}"
)
return None, 0
else:
logger.critical(
f"模型 {self.model_name} 请求因未知错误失败: {str(exception.__class__.__name__)} - {str(exception)}"
)
current_key_placeholder = request_content.get("current_key", "******")
handled_payload = await _safely_record(request_content, payload)
logger.critical(
f"请求头: {await self._build_headers(api_key=current_key_placeholder, no_key=True)} 请求体: {handled_payload}"
)
raise RuntimeError(f"模型 {self.model_name} API请求失败: {str(exception)}") from exception
async def _transform_parameters(self, merged_params: dict) -> dict:
"""根据模型名称转换合并后的参数，并移除内部参数"""
# (代码不变)
new_params = dict(merged_params)
new_params.pop("request_type", None)
if not self.is_gemini and self.model_name.lower() in self.MODELS_NEEDING_TRANSFORMATION:
new_params.pop("temperature", None)
if "max_tokens" in new_params:
new_params["max_completion_tokens"] = new_params.pop("max_tokens")
elif self.is_gemini:
gen_config = new_params.get("generationConfig", {})
if "temperature" in new_params:
gen_config["temperature"] = new_params.pop("temperature")
if "max_tokens" in new_params:
gen_config["maxOutputTokens"] = new_params.pop("max_tokens")
if "top_p" in new_params:
gen_config["topP"] = new_params.pop("top_p")
if "top_k" in new_params:
gen_config["topK"] = new_params.pop("top_k")
if gen_config:
new_params["generationConfig"] = gen_config
new_params.pop("frequency_penalty", None)
new_params.pop("presence_penalty", None)
new_params.pop("max_completion_tokens", None)
return new_params
async def _build_payload(
self, prompt: str, image_base64: str = None, image_format: str = None, merged_params: dict = None
) -> dict:
"""构建请求体 (区分 Gemini 和 OpenAI)，使用合并和转换后的参数"""
# (代码不变)
if merged_params is None:
merged_params = self.params
params_copy = await self._transform_parameters(merged_params)
if self.is_gemini:
parts = []
if prompt:
parts.append({"text": prompt})
if image_base64:
mime_type = f"image/{image_format.lower() if image_format else 'jpeg'}"
parts.append({"inlineData": {"mimeType": mime_type, "data": image_base64}})
payload = {"contents": [{"parts": parts}], **params_copy}
payload.pop("model", None)
# --- 添加 Gemini 安全设置 ---
safety_settings = [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"},
]
payload["safetySettings"] = safety_settings
logger.debug(f"模型 {self.model_name}: 已为 Gemini 函数调用请求添加 safetySettings (BLOCK_NONE)。")
# --- 结束添加安全设置 ---
else:
if image_base64:
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:image/{image_format.lower() if image_format else 'jpeg'};base64,{image_base64}"
},
},
],
}
]
else:
messages = [{"role": "user", "content": prompt}]
payload = {
"model": self.model_name,
"messages": messages,
**params_copy,
}
if "max_tokens" not in payload and "max_completion_tokens" not in payload:
if "max_tokens" not in params_copy and "max_completion_tokens" not in params_copy:
payload["max_tokens"] = global_config.model_max_output_length
if "max_completion_tokens" in payload:
payload["max_tokens"] = payload.pop("max_completion_tokens")
return payload
def _default_response_handler(
self, result: dict, user_id: str = "system", request_type: str = None, endpoint: str = "/chat/completions"
) -> Tuple:
"""默认响应解析 (区分 Gemini 和 OpenAI)，并处理函数/工具调用"""
content = "没有返回结果"
reasoning_content = ""
tool_calls = None # OpenAI 格式
function_call = None # Gemini 格式
prompt_tokens = 0
completion_tokens = 0
total_tokens = 0
if self.is_gemini:
# --- 解析 Gemini 响应 ---
try:
if "candidates" in result and result["candidates"]:
candidate = result["candidates"][0]
# 检查是否有 content 和 parts
if "content" in candidate and "parts" in candidate["content"] and candidate["content"]["parts"]:
# 查找 functionCall 或 text 部分
final_text_parts = []
for part in candidate["content"]["parts"]:
if "functionCall" in part:
function_call = part["functionCall"] # 获取 Gemini 的 functionCall
# Gemini functionCall 通常不与 text 一起返回，这里假设只处理 functionCall
break # 找到 functionCall 就停止处理 parts
elif "text" in part:
final_text_parts.append(part.get("text", ""))
if not function_call: # 如果没有 functionCall，处理 text
raw_content = "".join(final_text_parts).strip()
content, reasoning = self._extract_reasoning(raw_content)
reasoning_content = reasoning
# else: function_call 已获取，content 留空或设为特定值
else:
content = "Gemini响应中缺少 content 或 parts"
logger.warning(f"模型 {self.model_name}: Gemini 响应格式不完整 (缺少 content/parts): {result}")
finish_reason = candidate.get("finishReason")
if finish_reason == "SAFETY":
logger.warning(f"模型 {self.model_name}: Gemini 响应因安全设置被阻止。")
content = "响应内容因安全原因被过滤。"
elif finish_reason == "RECITATION":
logger.warning(f"模型 {self.model_name}: Gemini 响应因引用限制被阻止。")
content = "响应内容因引用限制被过滤。"
elif finish_reason == "OTHER":
logger.warning(f"模型 {self.model_name}: Gemini 响应因未知原因停止。")
# finishReason == "TOOL_CODE" or "FUNCTION_CALL" 是正常情况
usage = result.get("usageMetadata", {})
if usage:
prompt_tokens = usage.get("promptTokenCount", 0)
completion_tokens = usage.get("candidatesTokenCount", 0)
total_tokens = usage.get("totalTokenCount", 0)
if completion_tokens == 0 and total_tokens > 0:
completion_tokens = total_tokens - prompt_tokens
else:
logger.warning(f"模型 {self.model_name} (Gemini) 的响应中缺少 'usageMetadata' 信息。")
except Exception as e:
logger.error(f"解析 Gemini 响应出错: {e} - 响应: {result}")
content = "解析 Gemini 响应时出错"
else:
# --- 解析 OpenAI 兼容响应 ---
# (代码不变)
if "choices" in result and result["choices"]:
message = result["choices"][0].get("message", {})
raw_content = message.get("content", "")
content, reasoning = self._extract_reasoning(raw_content if raw_content else "")
explicit_reasoning = message.get("model_extra", {}).get("reasoning_content", "")
if not explicit_reasoning:
explicit_reasoning = message.get("reasoning_content", "")
reasoning_content = explicit_reasoning if explicit_reasoning else reasoning
tool_calls = message.get("tool_calls", None) # 获取 OpenAI 的 tool_calls
usage = result.get("usage", {})
if usage:
prompt_tokens = usage.get("prompt_tokens", 0)
completion_tokens = usage.get("completion_tokens", 0)
total_tokens = usage.get("total_tokens", 0)
else:
logger.warning(f"模型 {self.model_name} (OpenAI) 的响应中缺少 'usage' 信息。")
else:
logger.warning(f"模型 {self.model_name} (OpenAI) 的响应格式不符合预期: {result}")
# --- 记录 Token 使用情况 ---
# (代码不变)
if prompt_tokens > 0 or completion_tokens > 0 or total_tokens > 0:
self._record_usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
user_id=user_id,
request_type=request_type,
endpoint=endpoint,
)
else:
logger.warning(f"模型 {self.model_name}: 未能从响应中提取有效的 token 使用信息。")
# --- 返回结果 (统一格式) ---
final_tool_calls = None
if tool_calls: # 来自 OpenAI
final_tool_calls = tool_calls
logger.debug(f"检测到 OpenAI 工具调用: {final_tool_calls}")
elif function_call: # 来自 Gemini
logger.debug(f"检测到 Gemini 函数调用: {function_call}")
# 将 Gemini functionCall 转换为 OpenAI tool_calls 格式
# 注意: Gemini 的 functionCall 没有显式的 id 和 type，需要模拟
final_tool_calls = [
{
"id": f"call_{random.randint(1000, 9999)}", # 生成一个随机 ID
"type": "function",
"function": {
"name": function_call.get("name"),
# Gemini 的参数在 'args' 中，OpenAI 在 'arguments' (通常是 JSON 字符串)
# 需要将 Gemini 的 dict 参数转换为 JSON 字符串
"arguments": json.dumps(function_call.get("args", {})),
},
}
]
logger.debug(f"转换为 OpenAI tool_calls 格式: {final_tool_calls}")
if final_tool_calls:
# 如果有工具/函数调用，通常 content 为空或包含思考过程，这里返回转换后的调用信息
return content, reasoning_content, final_tool_calls
else:
# 没有工具/函数调用，返回普通文本响应
return content, reasoning_content
@staticmethod
def _extract_reasoning(content: str) -> Tuple[str, str]:
"""CoT思维链提取"""
# (代码不变)
if not content:
return "", ""
match = re.search(r"(.*?)", content, re.DOTALL)
cleaned_content = re.sub(r".*?", "", content, flags=re.DOTALL, count=1).strip()
if match:
reasoning = match.group(1).strip()
else:
reasoning = ""
return cleaned_content, reasoning
async def _build_headers(self, api_key: str, no_key: bool = False) -> dict:
"""构建请求头 (区分 Gemini 和 OpenAI)"""
# (代码不变)
if no_key:
if self.is_gemini:
return {"x-goog-api-key": "**********", "Content-Type": "application/json"}
else:
return {"Authorization": "Bearer **********", "Content-Type": "application/json"}
else:
if not api_key:
logger.error(f"尝试使用无效 (空) 的 API key 为模型 {self.model_name} 构建请求头。")
raise ValueError("无效的 API key 提供给 _build_headers。")
if self.is_gemini:
return {"x-goog-api-key": api_key, "Content-Type": "application/json"}
else:
return {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
async def generate_response(self, prompt: str, user_id: str = "system", **kwargs) -> Tuple:
"""根据输入的提示生成模型的异步响应，支持覆盖参数"""
endpoint = ":generateContent" if self.is_gemini else "/chat/completions"
response = await self._execute_request(
endpoint=endpoint, prompt=prompt, user_id=user_id, request_type="chat", **kwargs
)
if len(response) == 3:
content, reasoning_content, tool_calls = response
return content, reasoning_content, self.model_name, tool_calls
else:
content, reasoning_content = response
return content, reasoning_content, self.model_name
async def generate_response_for_image(
self, prompt: str, image_base64: str, image_format: str, user_id: str = "system", **kwargs
) -> Tuple:
"""根据输入的提示和图片生成模型的异步响应，支持覆盖参数"""
endpoint = ":generateContent" if self.is_gemini else "/chat/completions"
response = await self._execute_request(
endpoint=endpoint,
prompt=prompt,
image_base64=image_base64,
image_format=image_format,
user_id=user_id,
request_type="vision",
**kwargs,
)
# _default_response_handler 现在总是返回至少2个值
if len(response) == 3:
return response # content, reasoning, tool_calls (tool_calls 可能为 None)
elif len(response) == 2:
content, reasoning = response
return content, reasoning # 对于 vision 请求，通常没有 tool_calls
else:
logger.error(f"来自 _default_response_handler 的意外响应格式: {response}")
return "处理响应出错", ""
async def generate_response_async(
self, prompt: str, user_id: str = "system", request_type: str = "chat", **kwargs
) -> Union[str, Tuple]:
"""异步方式根据输入的提示生成模型的响应 (通用)，支持覆盖参数"""
# (代码不变)
endpoint = ":generateContent" if self.is_gemini else "/chat/completions"
response = await self._execute_request(
endpoint=endpoint,
prompt=prompt,
payload=None,
retry_policy=None,
response_handler=None,
user_id=user_id,
request_type=request_type,
**kwargs,
)
return response
# 修改：实现 Gemini Function Calling 的 Payload 构建
async def generate_response_tool_async(
self, prompt: str, tools: list, user_id: str = "system", **kwargs
) -> tuple[str, str, list | None]:
"""异步方式根据输入的提示和工具生成模型的响应，支持覆盖参数和 Gemini 函数调用"""
endpoint = ":generateContent" if self.is_gemini else "/chat/completions"
merged_params = {**self.params, **kwargs}
transformed_params = await self._transform_parameters(merged_params) # 清理 request_type 等
payload = None
if self.is_gemini:
# --- 构建 Gemini Function Calling Payload ---
logger.debug(f"为 Gemini ({self.model_name}) 构建函数调用请求。")
# 1. 转换工具定义 (OpenAI -> Gemini)
# OpenAI tool format: [{"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}}]
# Gemini tool format: [{"functionDeclarations": [{"name": ..., "description": ..., "parameters": ...}]}]
function_declarations = []
if tools:
for tool in tools:
if tool.get("type") == "function" and "function" in tool:
func_def = tool["function"]
# Gemini parameters 使用 OpenAPI Schema，与 OpenAI 基本兼容
function_declarations.append(
{
"name": func_def.get("name"),
"description": func_def.get("description", ""), # Description is required for Gemini
"parameters": func_def.get(
"parameters", {"type": "object", "properties": {}}
), # Ensure parameters exist
}
)
else:
logger.warning(f"跳过不支持的工具类型或格式: {tool}")
if not function_declarations:
logger.error("没有有效的函数声明可用于 Gemini 请求。")
return "没有提供有效的函数定义", "", None
gemini_tools = [{"functionDeclarations": function_declarations}]
# 2. 构建 Gemini Payload
# parts = [{"text": prompt}] # 初始 parts
payload = {
"contents": [{"parts": [{"text": prompt}]}], # 包含用户提示
"tools": gemini_tools,
# toolConfig 默认是 AUTO，可以根据需要从 kwargs 获取或硬编码
# "toolConfig": {"functionCallingConfig": {"mode": "ANY"}}, # 例如强制调用
**transformed_params, # 合并其他转换后的参数 (如 generationConfig)
}
payload.pop("model", None) # Gemini 不在顶层传 model
payload.pop("messages", None) # 移除 OpenAI 特有的 messages
payload.pop("tool_choice", None) # 移除 OpenAI 特有的 tool_choice
logger.trace(f"构建的 Gemini 函数调用 Payload: {json.dumps(payload, indent=2)}")
else:
# --- 构建 OpenAI Tool Calling Payload ---
# (逻辑不变)
logger.debug(f"为 OpenAI 兼容模型 ({self.model_name}) 构建工具调用请求。")
payload = {
"model": self.model_name,
"messages": [{"role": "user", "content": prompt}],
**transformed_params,
"tools": tools,
"tool_choice": transformed_params.get("tool_choice", "auto"),
}
if "max_completion_tokens" in payload:
payload["max_tokens"] = payload.pop("max_completion_tokens")
if "max_tokens" not in payload:
payload["max_tokens"] = global_config.model_max_output_length
# --- 执行请求 ---
if payload is None:
logger.error("未能构建有效的 API 请求 payload。")
return "内部错误：无法构建请求", "", None
response = await self._execute_request(
endpoint=endpoint,
payload=payload,
prompt=prompt, # prompt 仍然需要，用于可能的重试
user_id=user_id,
request_type="tool_call",
**kwargs, # 传递原始 kwargs 以便在重试时重新合并
)
# _default_response_handler 现在会处理 Gemini functionCall 并统一格式
logger.debug(f"模型 {self.model_name} 工具/函数调用返回结果: {response}")
if isinstance(response, tuple) and len(response) == 3:
content, reasoning_content, final_tool_calls = response
# final_tool_calls 已经是统一的 OpenAI 格式
return content, reasoning_content, final_tool_calls
elif isinstance(response, tuple) and len(response) == 2:
content, reasoning_content = response
logger.debug("收到普通响应，无工具/函数调用")
return content, reasoning_content, None
else:
logger.error(f"收到来自 _execute_request/_default_response_handler 的意外响应格式: {response}")
return "处理响应时出错", "", None
async def get_embedding(self, text: str, user_id: str = "system", **kwargs) -> Union[list, None]:
"""异步方法：获取文本的embedding向量，支持覆盖参数 (Gemini Embedding 需注意模型名称)"""
# (代码不变)
if len(text) < 1:
logger.debug("该消息没有长度，不再发送获取embedding向量的请求")
return None
api_kwargs = {k: v for k, v in kwargs.items() if k != "request_type"}
if self.is_gemini:
endpoint = ":embedContent"
payload = {"model": f"models/{self.model_name}", "content": {"parts": [{"text": text}]}, **api_kwargs}
payload.pop("encoding_format", None)
payload.pop("input", None)
else:
endpoint = "/embeddings"
payload = {"model": self.model_name, "input": text, "encoding_format": "float", **api_kwargs}
payload.pop("content", None)
payload.pop("taskType", None)
def embedding_handler(result):
# (代码不变)
embedding_value = None
prompt_tokens = 0
completion_tokens = 0
total_tokens = 0
if self.is_gemini:
if "embedding" in result and "value" in result["embedding"]:
embedding_value = result["embedding"]["value"]
logger.warning(f"模型 {self.model_name} (Gemini Embedding): 响应中未找到明确的 token 使用信息。")
else:
if "data" in result and len(result["data"]) > 0:
embedding_value = result["data"][0].get("embedding", None)
usage = result.get("usage", {})
if usage:
prompt_tokens = usage.get("prompt_tokens", 0)
total_tokens = usage.get("total_tokens", 0)
else:
logger.warning(f"模型 {self.model_name} (OpenAI Embedding) 的响应中缺少 'usage' 信息。")
self._record_usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
user_id=user_id,
request_type="embedding",
endpoint=endpoint,
)
return embedding_value
embedding = await self._execute_request(
endpoint=endpoint,
payload=payload,
prompt=text,
retry_policy={"max_retries": 2, "base_wait": 6},
response_handler=embedding_handler,
user_id=user_id,
request_type="embedding",
**api_kwargs,
)
return embedding
def compress_base64_image_by_scale(base64_data: str, target_size: int = 0.8 * 1024 * 1024) -> str:
"""压缩base64格式的图片到指定大小"""
# (代码不变)
try:
image_data = base64.b64decode(base64_data)
if len(image_data) <= target_size * 1.05:
logger.info(f"图片大小 {len(image_data) / 1024:.1f}KB 已足够小，无需压缩。")
return base64_data
img = Image.open(io.BytesIO(image_data))
img_format = img.format
original_width, original_height = img.size
scale = max(0.2, min(1.0, (target_size / len(image_data)) ** 0.5))
new_width = max(1, int(original_width * scale))
new_height = max(1, int(original_height * scale))
output_buffer = io.BytesIO()
save_format = img_format # Default to original format
if getattr(img, "is_animated", False) and img.n_frames > 1:
frames = []
durations = []
loop = img.info.get("loop", 0)
disposal = img.info.get("disposal", 2)
logger.info(f"检测到 GIF 动图 ({img.n_frames} 帧)，尝试按比例压缩...")
for frame_idx in range(img.n_frames):
img.seek(frame_idx)
current_duration = img.info.get("duration", 100)
durations.append(current_duration)
new_frame = img.convert("RGBA").copy()
resized_frame = new_frame.resize((new_width, new_height), Image.Resampling.LANCZOS)
frames.append(resized_frame)
if frames:
frames[0].save(
output_buffer,
format="GIF",
save_all=True,
append_images=frames[1:],
optimize=False,
duration=durations,
loop=loop,
disposal=disposal,
transparency=img.info.get("transparency", None),
background=img.info.get("background", None),
)
save_format = "GIF"
else:
logger.warning("未能处理 GIF 帧。")
return base64_data
else:
if img.mode in ("RGBA", "LA") or "transparency" in img.info:
resized_img = img.convert("RGBA").resize((new_width, new_height), Image.Resampling.LANCZOS)
save_format = "PNG"
save_params = {"optimize": True}
else:
resized_img = img.convert("RGB").resize((new_width, new_height), Image.Resampling.LANCZOS)
if img_format and img_format.upper() == "JPEG":
save_format = "JPEG"
save_params = {"quality": 85, "optimize": True}
else:
save_format = "PNG"
save_params = {"optimize": True}
resized_img.save(output_buffer, format=save_format, **save_params)
compressed_data = output_buffer.getvalue()
logger.success(
f"压缩图片: {original_width}x{original_height} -> {new_width}x{new_height} ({img.format} -> {save_format})"
)
logger.info(
f"压缩前大小: {len(image_data) / 1024:.1f}KB, 压缩后大小: {len(compressed_data) / 1024:.1f}KB (目标: {target_size / 1024:.1f}KB)"
)
if len(compressed_data) < len(image_data) * 0.95:
return base64.b64encode(compressed_data).decode("utf-8")
else:
logger.info("压缩效果不明显或反而增大，返回原始图片。")
return base64_data
except Exception as e:
logger.error(f"压缩图片失败: {str(e)}")
import traceback
logger.error(traceback.format_exc())
return base64_data