fix(gemini): 对 thinking_budget 不同模型的处理

pull/1203/head
foxplaying 2025-08-21 03:01:25 +08:00 committed by GitHub
parent a68c68cbe9
commit 6ab8edd566
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 37 additions and 10 deletions

View File

@ -44,10 +44,13 @@ from ..payload_content.tool_option import ToolOption, ToolParam, ToolCall
logger = get_logger("Gemini客户端") logger = get_logger("Gemini客户端")
# gemini_thinking参数 # gemini_thinking参数默认范围
GEMINI_THINKING_BUDGET_MIN = 512 # 对不同模型的思考预算范围配置
GEMINI_THINKING_BUDGET_MAX = 24576 THINKING_BUDGET_LIMITS = {
DEFAULT_THINKING_BUDGET = 1024 "gemini-2.5-flash": {"min": 1, "max": 24576, "can_disable": True},
"gemini-2.5-flash-lite": {"min": 512, "max": 24576, "can_disable": True},
"gemini-2.5-pro": {"min": 128, "max": 32768, "can_disable": False},
}
gemini_safe_settings = [ gemini_safe_settings = [
SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE), SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE),
@ -333,6 +336,30 @@ class GeminiClient(BaseClient):
api_key=api_provider.api_key, api_key=api_provider.api_key,
) # 这里和openai不一样gemini会自己决定自己是否需要retry ) # 这里和openai不一样gemini会自己决定自己是否需要retry
# 思维预算的特殊值
TB_DYNAMIC_MODE = -1
TB_DISABLE_OR_MIN = 0
@staticmethod
def clamp_thinking_budget(tb: int, model_id: str) -> int:
"""
按模型限制思考预算范围仅支持指定的模型
"""
limits = THINKING_BUDGET_LIMITS.get(model_id)
if limits is None:
raise ValueError(f"模型 {model_id} 不支持 ThinkingConfig")
if tb == GeminiClient.TB_DYNAMIC_MODE:
return GeminiClient.TB_DYNAMIC_MODE # 动态思考模式
if tb == GeminiClient.TB_DISABLE_OR_MIN:
if limits["can_disable"]:
# 允许禁用思考预算
return GeminiClient.TB_DISABLE_OR_MIN
else:
# 不允许禁用,返回最小值
return limits["min"]
# 正常范围裁剪
return max(limits["min"], min(tb, limits["max"]))
async def get_response( async def get_response(
self, self,
model_info: ModelInfo, model_info: ModelInfo,
@ -379,17 +406,17 @@ class GeminiClient(BaseClient):
# 将tool_options转换为Gemini API所需的格式 # 将tool_options转换为Gemini API所需的格式
tools = _convert_tool_options(tool_options) if tool_options else None tools = _convert_tool_options(tool_options) if tool_options else None
# 将response_format转换为Gemini API所需的格式 # 将response_format转换为Gemini API所需的格式
# 处理 thinking_budget
try: try:
if extra_params and "thinking_budget" in extra_params: if extra_params and "thinking_budget" in extra_params:
tb = extra_params["thinking_budget"] tb = int(extra_params["thinking_budget"])
tb = int(tb) # 尝试转换为整数
else: else:
tb = int(max_tokens / 2) tb = int(max_tokens / 2)
except (TypeError, ValueError) as e: except (ValueError, TypeError):
logger.warning(f"无效的thinking_budget值 {extra_params.get('thinking_budget') if extra_params else None}使用默认值 {DEFAULT_THINKING_BUDGET}: {e}") logger.warning(f"无效的thinking_budget值 {extra_params.get('thinking_budget') if extra_params else None}将使用 max_tokens/2 作为默认值")
tb = DEFAULT_THINKING_BUDGET tb = int(max_tokens / 2)
tb = max(GEMINI_THINKING_BUDGET_MIN, min(tb, GEMINI_THINKING_BUDGET_MAX)) # 限制在合法范围512-24576 tb = self.clamp_thinking_budget(tb, model_info.model_identifier)
generation_config_dict = { generation_config_dict = {
"max_output_tokens": max_tokens, "max_output_tokens": max_tokens,