Skip to content

Commit 85a5d34

Browse files
authored
Merge pull request #189 from NyxJae/dev
feat: 在env中添加URL Context功能默认开关,思考预算和Google搜索功能的默认配置选项
2 parents 833ec19 + 8c4a97c commit 85a5d34

4 files changed

Lines changed: 154 additions & 93 deletions

File tree

.env.example

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,21 @@ DEFAULT_TOP_P=0.95
8585
# 默认停止序列 (JSON 数组格式)
8686
DEFAULT_STOP_SEQUENCES=["用户:"]
8787

88+
# 是否在处理请求时自动打开并使用 "URL Context" 功能,此工具功能详情可参考:https://ai.google.dev/gemini-api/docs/url-context
89+
ENABLE_URL_CONTEXT=true
90+
91+
# 是否默认启用 "指定思考预算" 功能 (true/false),不启用时模型一般将自行决定思考预算
92+
# 当 API 请求中未提供 reasoning_effort 参数时将使用此值。
93+
ENABLE_THINKING_BUDGET=false
94+
95+
# "指定思考预算量" 的默认值 (token)
96+
# 当 API 请求中未提供 reasoning_effort 参数时,将使用此值。
97+
DEFAULT_THINKING_BUDGET=8192
98+
99+
# 是否默认启用 "Google Search" 功能 (true/false)
100+
# 当 API 请求中未提供 tools 参数时,将使用此设置作为 Google Search 的默认开关状态。
101+
ENABLE_GOOGLE_SEARCH=false
102+
88103
# =============================================================================
89104
# 超时配置 (毫秒)
90105
# =============================================================================

browser_utils/page_controller.py

Lines changed: 118 additions & 93 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
)
1919
from config import (
2020
CLICK_TIMEOUT_MS, WAIT_FOR_ELEMENT_TIMEOUT_MS, CLEAR_CHAT_VERIFY_TIMEOUT_MS,
21-
DEFAULT_TEMPERATURE, DEFAULT_MAX_OUTPUT_TOKENS, DEFAULT_STOP_SEQUENCES, DEFAULT_TOP_P
21+
DEFAULT_TEMPERATURE, DEFAULT_MAX_OUTPUT_TOKENS, DEFAULT_STOP_SEQUENCES, DEFAULT_TOP_P,
22+
ENABLE_URL_CONTEXT, ENABLE_THINKING_BUDGET, DEFAULT_THINKING_BUDGET, ENABLE_GOOGLE_SEARCH
2223
)
2324
from models import ClientDisconnectedError
2425
from .operations import save_error_snapshot, _wait_for_response_completion, _get_final_response_content
@@ -62,36 +63,70 @@ async def adjust_parameters(self, request_params: Dict[str, Any], page_params_ca
6263
await self._check_disconnect(check_client_disconnected, "End Parameter Adjustment")
6364

6465
# 调整URL CONTEXT
65-
await self._open_url_content(check_client_disconnected)
66-
67-
# 调整“思考预算”开关
68-
await self._set_thinking_budget_toggle_checked(check_client_disconnected)
66+
if ENABLE_URL_CONTEXT:
67+
await self._open_url_content(check_client_disconnected)
68+
else:
69+
self.logger.info(f"[{self.req_id}] URL Context 功能已禁用,跳过调整。")
6970

70-
# 在函数末尾,await self._set_thinking_budget_toggle_checked(...) 之后添加:
71-
self.logger.info(f"[{self.req_id}] 检查预算...{request_params.get('reasoning_effort')}")
72-
if 'reasoning_effort' in request_params:
73-
await self._adjust_thinking_budget(request_params.get('reasoning_effort'), check_client_disconnected)
71+
# 调整“思考预算”
72+
await self._handle_thinking_budget(request_params, check_client_disconnected)
7473

7574
# 调整 Google Search 开关
7675
await self._adjust_google_search(request_params, check_client_disconnected)
7776

78-
async def _adjust_thinking_budget(self, reasoning_effort: Optional[str], check_client_disconnected: Callable):
77+
async def _handle_thinking_budget(self, request_params: Dict[str, Any], check_client_disconnected: Callable):
78+
"""处理思考预算的调整逻辑。"""
79+
reasoning_effort = request_params.get('reasoning_effort')
80+
if reasoning_effort is not None:
81+
# 用户指定了,则开启并设置
82+
self.logger.info(f"[{self.req_id}] 用户指定了 reasoning_effort: {reasoning_effort}。")
83+
await self._control_thinking_budget_toggle(should_be_checked=True, check_client_disconnected=check_client_disconnected)
84+
await self._adjust_thinking_budget(reasoning_effort, check_client_disconnected)
85+
else:
86+
# 用户未指定,根据默认配置
87+
self.logger.info(f"[{self.req_id}] 用户未指定 reasoning_effort,根据默认配置 ENABLE_THINKING_BUDGET: {ENABLE_THINKING_BUDGET}。")
88+
await self._control_thinking_budget_toggle(should_be_checked=ENABLE_THINKING_BUDGET, check_client_disconnected=check_client_disconnected)
89+
if ENABLE_THINKING_BUDGET:
90+
# 如果默认开启,则使用默认值
91+
await self._adjust_thinking_budget(None, check_client_disconnected)
92+
93+
def _parse_thinking_budget(self, reasoning_effort: Optional[Any]) -> Optional[int]:
94+
"""从 reasoning_effort 解析出 token_budget。"""
95+
token_budget = None
96+
if reasoning_effort is None:
97+
token_budget = DEFAULT_THINKING_BUDGET
98+
self.logger.info(f"[{self.req_id}] 'reasoning_effort' 为空,使用默认思考预算: {token_budget}")
99+
elif isinstance(reasoning_effort, int):
100+
token_budget = reasoning_effort
101+
elif isinstance(reasoning_effort, str):
102+
if reasoning_effort.lower() == 'none':
103+
token_budget = DEFAULT_THINKING_BUDGET
104+
self.logger.info(f"[{self.req_id}] 'reasoning_effort' 为 'none' 字符串,使用默认思考预算: {token_budget}")
105+
else:
106+
effort_map = {
107+
"low": 1000,
108+
"medium": 8000,
109+
"high": 24000
110+
}
111+
token_budget = effort_map.get(reasoning_effort.lower())
112+
if token_budget is None:
113+
try:
114+
token_budget = int(reasoning_effort)
115+
except (ValueError, TypeError):
116+
pass # token_budget remains None
117+
118+
if token_budget is None:
119+
self.logger.warning(f"[{self.req_id}] 无法从 '{reasoning_effort}' (类型: {type(reasoning_effort)}) 解析出有效的 token_budget。")
120+
121+
return token_budget
122+
123+
async def _adjust_thinking_budget(self, reasoning_effort: Optional[Any], check_client_disconnected: Callable):
79124
"""根据 reasoning_effort 调整思考预算。"""
80-
self.logger.info(f"[{self.req_id}] 检查并调整思考预算...")
125+
self.logger.info(f"[{self.req_id}] 检查并调整思考预算,输入值: {reasoning_effort}")
81126

82-
token_budget = None
83-
if reasoning_effort is None or reasoning_effort.lower() == 'none':
84-
token_budget = 8192
85-
self.logger.info(f"[{self.req_id}] 'reasoning_effort' 为空或 'none',使用默认思考预算: {token_budget}")
86-
else:
87-
effort_map = {
88-
"low": 1000,
89-
"medium": 8000,
90-
"high": 24000
91-
}
92-
token_budget = effort_map.get(reasoning_effort.lower())
93-
94-
if not token_budget:
127+
token_budget = self._parse_thinking_budget(reasoning_effort)
128+
129+
if token_budget is None:
95130
self.logger.warning(f"[{self.req_id}] 无效的 reasoning_effort 值: '{reasoning_effort}'。跳过调整。")
96131
return
97132

@@ -118,18 +153,32 @@ async def _adjust_thinking_budget(self, reasoning_effort: Optional[str], check_c
118153
if isinstance(e, ClientDisconnectedError):
119154
raise
120155

156+
def _should_enable_google_search(self, request_params: Dict[str, Any]) -> bool:
157+
"""根据请求参数或默认配置决定是否应启用 Google Search。"""
158+
if 'tools' in request_params and request_params.get('tools') is not None:
159+
tools = request_params.get('tools')
160+
has_google_search_tool = False
161+
if isinstance(tools, list):
162+
for tool in tools:
163+
if isinstance(tool, dict):
164+
if tool.get('google_search_retrieval') is not None:
165+
has_google_search_tool = True
166+
break
167+
if tool.get('function', {}).get('name') == 'googleSearch':
168+
has_google_search_tool = True
169+
break
170+
self.logger.info(f"[{self.req_id}] 请求中包含 'tools' 参数。检测到 Google Search 工具: {has_google_search_tool}。")
171+
return has_google_search_tool
172+
else:
173+
self.logger.info(f"[{self.req_id}] 请求中不包含 'tools' 参数。使用默认配置 ENABLE_GOOGLE_SEARCH: {ENABLE_GOOGLE_SEARCH}。")
174+
return ENABLE_GOOGLE_SEARCH
175+
121176
async def _adjust_google_search(self, request_params: Dict[str, Any], check_client_disconnected: Callable):
122-
"""根据请求参数中的 'googleSearch' 工具存在与否,双向控制 Google Search 开关。"""
177+
"""根据请求参数或默认配置,双向控制 Google Search 开关。"""
123178
self.logger.info(f"[{self.req_id}] 检查并调整 Google Search 开关...")
124179

125-
tools = request_params.get('tools')
126-
has_google_search = False
127-
if isinstance(tools, list):
128-
for tool in tools:
129-
if isinstance(tool, dict) and tool.get('function', {}).get('name') == 'googleSearch':
130-
has_google_search = True
131-
break
132-
180+
should_enable_search = self._should_enable_google_search(request_params)
181+
133182
toggle_selector = GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR
134183

135184
try:
@@ -138,39 +187,25 @@ async def _adjust_google_search(self, request_params: Dict[str, Any], check_clie
138187
await self._check_disconnect(check_client_disconnected, "Google Search 开关 - 元素可见后")
139188

140189
is_checked_str = await toggle_locator.get_attribute("aria-checked")
141-
self.logger.info(f"[{self.req_id}] Google Search 开关 'aria-checked' 状态: '{is_checked_str}'。是否需要搜索: {has_google_search}")
142-
143-
if has_google_search:
144-
# 需要搜索,但开关是关闭的,则打开它
145-
if is_checked_str == "false":
146-
self.logger.info(f"[{self.req_id}] 'googleSearch' 工具存在,但开关关闭。正在点击以打开...")
147-
await toggle_locator.click(timeout=CLICK_TIMEOUT_MS)
148-
await self._check_disconnect(check_client_disconnected, "Google Search 开关 - 点击打开后")
149-
await asyncio.sleep(0.5) # 等待UI更新
150-
new_state = await toggle_locator.get_attribute("aria-checked")
151-
if new_state == "true":
152-
self.logger.info(f"[{self.req_id}] ✅ Google Search 开关已成功打开。")
153-
else:
154-
self.logger.warning(f"[{self.req_id}] ⚠️ Google Search 开关打开失败。当前状态: '{new_state}'")
190+
is_currently_checked = is_checked_str == "true"
191+
self.logger.info(f"[{self.req_id}] Google Search 开关当前状态: '{is_checked_str}'。期望状态: {should_enable_search}")
192+
193+
if should_enable_search != is_currently_checked:
194+
action = "打开" if should_enable_search else "关闭"
195+
self.logger.info(f"[{self.req_id}] Google Search 开关状态与期望不符。正在点击以{action}...")
196+
await toggle_locator.click(timeout=CLICK_TIMEOUT_MS)
197+
await self._check_disconnect(check_client_disconnected, f"Google Search 开关 - 点击{action}后")
198+
await asyncio.sleep(0.5) # 等待UI更新
199+
new_state = await toggle_locator.get_attribute("aria-checked")
200+
if (new_state == "true") == should_enable_search:
201+
self.logger.info(f"[{self.req_id}] ✅ Google Search 开关已成功{action}。")
155202
else:
156-
self.logger.info(f"[{self.req_id}] 'googleSearch' 工具存在,且开关已打开,无需操作。")
203+
self.logger.warning(f"[{self.req_id}] ⚠️ Google Search 开关{action}失败。当前状态: '{new_state}'")
157204
else:
158-
# 不需要搜索,但开关是打开的,则关闭它
159-
if is_checked_str == "true":
160-
self.logger.info(f"[{self.req_id}] 'googleSearch' 工具不存在,但开关打开。正在点击以关闭...")
161-
await toggle_locator.click(timeout=CLICK_TIMEOUT_MS)
162-
await self._check_disconnect(check_client_disconnected, "Google Search 开关 - 点击关闭后")
163-
await asyncio.sleep(0.5) # 等待UI更新
164-
new_state = await toggle_locator.get_attribute("aria-checked")
165-
if new_state == "false":
166-
self.logger.info(f"[{self.req_id}] ✅ Google Search 开关已成功关闭。")
167-
else:
168-
self.logger.warning(f"[{self.req_id}] ⚠️ Google Search 开关关闭失败。当前状态: '{new_state}'")
169-
else:
170-
self.logger.info(f"[{self.req_id}] 'googleSearch' 工具不存在,且开关已关闭,无需操作。")
205+
self.logger.info(f"[{self.req_id}] Google Search 开关已处于期望状态,无需操作。")
171206

172207
except Exception as e:
173-
self.logger.error(f"[{self.req_id}] ❌ 操作 'Google Search toggle' 开关 '{toggle_selector}' 时发生错误: {e}")
208+
self.logger.error(f"[{self.req_id}] ❌ 操作 'Google Search toggle' 开关时发生错误: {e}")
174209
if isinstance(e, ClientDisconnectedError):
175210
raise
176211

@@ -196,53 +231,43 @@ async def _open_url_content(self,check_client_disconnected: Callable):
196231
except Exception as e:
197232
self.logger.error(f"[{self.req_id}] ❌ 操作USE_URL_CONTEXT_SELECTOR时发生错误:{e}。")
198233

199-
async def _set_thinking_budget_toggle_checked(self, check_client_disconnected: Callable):
234+
async def _control_thinking_budget_toggle(self, should_be_checked: bool, check_client_disconnected: Callable):
200235
"""
201-
确保 "Thinking Budget" 滑块开关处于选中状态
236+
根据 should_be_checked 的值,控制 "Thinking Budget" 滑块开关的状态
202237
"""
203238
toggle_selector = SET_THINKING_BUDGET_TOGGLE_SELECTOR
204-
self.logger.info(f"[{self.req_id}] 检查并设置 'Thinking Budget toggle' 滑块开关 '{toggle_selector}' 的状态...")
239+
self.logger.info(f"[{self.req_id}] 控制 'Thinking Budget' 开关,期望状态: {'选中' if should_be_checked else '未选中'}...")
205240

206241
try:
207242
toggle_locator = self.page.locator(toggle_selector)
208-
209-
# 等待元素在 DOM 中可见
210243
await expect_async(toggle_locator).to_be_visible(timeout=5000)
211244
await self._check_disconnect(check_client_disconnected, "思考预算开关 - 元素可见后")
212245

213-
# 1. 检查当前 'aria-checked' 属性
214-
is_checked = await toggle_locator.get_attribute("aria-checked")
215-
self.logger.info(f"[{self.req_id}] 思考预算开关当前 'aria-checked' 状态: {is_checked}")
246+
is_checked_str = await toggle_locator.get_attribute("aria-checked")
247+
current_state_is_checked = is_checked_str == "true"
248+
self.logger.info(f"[{self.req_id}] 思考预算开关当前 'aria-checked' 状态: {is_checked_str} (当前是否选中: {current_state_is_checked})")
216249

217-
# 2. 如果开关未选中 (aria-checked="false"),则点击它
218-
if is_checked == "false":
219-
self.logger.info(f"[{self.req_id}] 思考预算开关当前未选中,正在点击以启用...")
250+
if current_state_is_checked != should_be_checked:
251+
action = "启用" if should_be_checked else "禁用"
252+
self.logger.info(f"[{self.req_id}] 思考预算开关当前状态与期望不符,正在点击以{action}...")
220253
await toggle_locator.click(timeout=CLICK_TIMEOUT_MS)
221-
await self._check_disconnect(check_client_disconnected, "思考预算开关 - 点击后")
222-
223-
# 3. 验证操作是否成功
224-
await asyncio.sleep(0.5) # 短暂等待,让 UI 状态更新
225-
new_state = await toggle_locator.get_attribute("aria-checked")
226-
if new_state == "true":
227-
self.logger.info(f"[{self.req_id}] ✅ 'Thinking Budget toggle' 已成功启用。新状态: {new_state}")
228-
else:
229-
self.logger.warning(f"[{self.req_id}] ⚠️ 'Thinking Budget toggle' 点击后验证失败。期望状态: 'true', 实际状态: '{new_state}'")
230-
# 如果需要,可以在此处添加错误快照
231-
# await save_error_snapshot(f"thinking_budget_toggle_verify_fail_{self.req_id}")
254+
await self._check_disconnect(check_client_disconnected, f"思考预算开关 - 点击{action}后")
255+
256+
await asyncio.sleep(0.5)
257+
new_state_str = await toggle_locator.get_attribute("aria-checked")
258+
new_state_is_checked = new_state_str == "true"
232259

233-
elif is_checked == "true":
234-
self.logger.info(f"[{self.req_id}] 'Thinking Budget toggle' 已处于选中状态,无需操作。")
260+
if new_state_is_checked == should_be_checked:
261+
self.logger.info(f"[{self.req_id}] ✅ 'Thinking Budget' 开关已成功{action}。新状态: {new_state_str}")
262+
else:
263+
self.logger.warning(f"[{self.req_id}] ⚠️ 'Thinking Budget' 开关{action}后验证失败。期望状态: '{should_be_checked}', 实际状态: '{new_state_str}'")
235264
else:
236-
# 处理 'aria-checked' 属性不存在或值无效的情况
237-
self.logger.warning(f"[{self.req_id}] 无法确定 'Thinking Budget toggle' 的状态,'aria-checked' 值为: '{is_checked}'。")
265+
self.logger.info(f"[{self.req_id}] 'Thinking Budget' 开关已处于期望状态,无需操作。")
238266

239267
except Exception as e:
240-
self.logger.error(f"[{self.req_id}] ❌ 操作 'Thinking Budget toggle' 滑块开关 '{toggle_selector}' 时发生错误: {e}")
241-
# 如果需要,可以在此处添加错误快照
242-
# await save_error_snapshot(f"thinking_budget_toggle_error_{self.req_id}")
243-
# 如果是客户端断开连接的特定错误,则重新抛出
268+
self.logger.error(f"[{self.req_id}] ❌ 操作 'Thinking Budget toggle' 开关时发生错误: {e}")
244269
if isinstance(e, ClientDisconnectedError):
245-
raise
270+
raise
246271
async def _adjust_temperature(self, temperature: float, page_params_cache: dict, params_cache_lock: asyncio.Lock, check_client_disconnected: Callable):
247272
"""调整温度参数。"""
248273
async with params_cache_lock:

config/constants.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@
1919
DEFAULT_TEMPERATURE = float(os.environ.get('DEFAULT_TEMPERATURE', '1.0'))
2020
DEFAULT_MAX_OUTPUT_TOKENS = int(os.environ.get('DEFAULT_MAX_OUTPUT_TOKENS', '65536'))
2121
DEFAULT_TOP_P = float(os.environ.get('DEFAULT_TOP_P', '0.95'))
22+
# --- 默认功能开关 ---
23+
ENABLE_URL_CONTEXT = os.environ.get('ENABLE_URL_CONTEXT', 'true').lower() in ('true', '1', 'yes')
24+
ENABLE_THINKING_BUDGET = os.environ.get('ENABLE_THINKING_BUDGET', 'false').lower() in ('true', '1', 'yes')
25+
DEFAULT_THINKING_BUDGET = int(os.environ.get('DEFAULT_THINKING_BUDGET', '8192'))
26+
ENABLE_GOOGLE_SEARCH = os.environ.get('ENABLE_GOOGLE_SEARCH', 'false').lower() in ('true', '1', 'yes')
2227

2328
# 默认停止序列 - 支持 JSON 格式配置
2429
try:

docs/environment-configuration.md

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,22 @@ DEFAULT_TOP_P=0.95
139139
140140
# 默认停止序列 (JSON 数组格式)
141141
DEFAULT_STOP_SEQUENCES=["用户:"]
142+
143+
# 是否在处理请求时自动打开并使用 "URL Context" 功能
144+
# 参考: https://ai.google.dev/gemini-api/docs/url-context
145+
ENABLE_URL_CONTEXT=true
146+
147+
# 是否默认启用 "指定思考预算" 功能 (true/false),不启用时模型一般将自行决定思考预算
148+
# 当 API 请求中未提供 reasoning_effort 参数时,将使用此值。
149+
ENABLE_THINKING_BUDGET=false
150+
151+
# "指定思考预算量" 的默认值 (token)
152+
# 当 API 请求中未提供 reasoning_effort 参数时,将使用此值。
153+
DEFAULT_THINKING_BUDGET=8192
154+
155+
# 是否默认启用 "Google Search" 功能 (true/false)
156+
# 当 API 请求中未提供 tools 参数时,将使用此设置作为 Google Search 的默认开关状态。
157+
ENABLE_GOOGLE_SEARCH=false
142158
```
143159

144160
### 超时配置

0 commit comments

Comments
 (0)