diff --git a/src/llm_models/model_client/gemini_client.py b/src/llm_models/model_client/gemini_client.py index c0d4588b..93a41a3d 100644 --- a/src/llm_models/model_client/gemini_client.py +++ b/src/llm_models/model_client/gemini_client.py @@ -44,10 +44,13 @@ from ..payload_content.tool_option import ToolOption, ToolParam, ToolCall logger = get_logger("Gemini客户端") -# gemini_thinking参数 -GEMINI_THINKING_BUDGET_MIN = 512 -GEMINI_THINKING_BUDGET_MAX = 24576 -DEFAULT_THINKING_BUDGET = 1024 +# gemini_thinking参数(默认范围) +# 不同模型的思考预算范围配置 +THINKING_BUDGET_LIMITS = { + "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 = [ SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE), @@ -333,6 +336,48 @@ class GeminiClient(BaseClient): api_key=api_provider.api_key, ) # 这里和openai不一样,gemini会自己决定自己是否需要retry + # 思维预算特殊值 + THINKING_BUDGET_AUTO = -1 # 自动调整思考预算,由模型决定 + THINKING_BUDGET_DISABLED = 0 # 禁用思考预算(如果模型允许禁用) + + @staticmethod + def clamp_thinking_budget(tb: int, model_id: str): + """ + 按模型限制思考预算范围,仅支持指定的模型(支持带数字后缀的新版本) + """ + limits = None + matched_key = None + + # 优先尝试精确匹配 + if model_id in THINKING_BUDGET_LIMITS: + limits = THINKING_BUDGET_LIMITS[model_id] + matched_key = model_id + else: + # 按 key 长度倒序,保证更长的(更具体的,如 -lite)优先 + sorted_keys = sorted(THINKING_BUDGET_LIMITS.keys(), key=len, reverse=True) + for key in sorted_keys: + # 必须满足:完全等于 或者 前缀匹配(带 "-" 边界) + if model_id == key or model_id.startswith(key + "-"): + limits = THINKING_BUDGET_LIMITS[key] + matched_key = key + break + + # 特殊值处理 + if tb == GeminiClient.THINKING_BUDGET_AUTO: + return GeminiClient.THINKING_BUDGET_AUTO + if tb == GeminiClient.THINKING_BUDGET_DISABLED: + if limits and limits.get("can_disable", False): + return GeminiClient.THINKING_BUDGET_DISABLED + return limits["min"] if limits else GeminiClient.THINKING_BUDGET_AUTO + + # 已知模型裁剪到范围 + if limits: + return max(limits["min"], min(tb, limits["max"])) + + # 未知模型,返回动态模式 + logger.warning(f"模型 {model_id} 未在 THINKING_BUDGET_LIMITS 中定义,将使用动态模式 tb=-1 兼容。") + return GeminiClient.THINKING_BUDGET_AUTO + async def get_response( self, model_info: ModelInfo, @@ -378,19 +423,20 @@ class GeminiClient(BaseClient): messages = _convert_messages(message_list) # 将tool_options转换为Gemini API所需的格式 tools = _convert_tool_options(tool_options) if tool_options else None + + tb = GeminiClient.THINKING_BUDGET_AUTO + #空处理 + if extra_params and "thinking_budget" in extra_params: + try: + tb = int(extra_params["thinking_budget"]) + except (ValueError, TypeError): + logger.warning( + f"无效的 thinking_budget 值 {extra_params['thinking_budget']},将使用默认动态模式 {tb}" + ) + # 裁剪到模型支持的范围 + tb = self.clamp_thinking_budget(tb, model_info.model_identifier) + # 将response_format转换为Gemini API所需的格式 - try: - if extra_params and "thinking_budget" in extra_params: - tb = extra_params["thinking_budget"] - tb = int(tb) # 尝试转换为整数 - else: - tb = int(max_tokens / 2) - except (TypeError, ValueError) as e: - logger.warning(f"无效的thinking_budget值 {extra_params.get('thinking_budget') if extra_params else None},使用默认值 {DEFAULT_THINKING_BUDGET}: {e}") - tb = DEFAULT_THINKING_BUDGET - - tb = max(GEMINI_THINKING_BUDGET_MIN, min(tb, GEMINI_THINKING_BUDGET_MAX)) # 限制在合法范围(512-24576) - generation_config_dict = { "max_output_tokens": max_tokens, "temperature": temperature,