33封装了所有与Playwright页面直接交互的复杂逻辑。
44"""
55import asyncio
6- from typing import Callable , List , Dict , Any
6+ from typing import Callable , List , Dict , Any , Optional
77
88from playwright .async_api import Page as AsyncPage , expect as expect_async , TimeoutError
99
1212 MAT_CHIP_REMOVE_BUTTON_SELECTOR , TOP_P_INPUT_SELECTOR , SUBMIT_BUTTON_SELECTOR ,
1313 CLEAR_CHAT_BUTTON_SELECTOR , CLEAR_CHAT_CONFIRM_BUTTON_SELECTOR , OVERLAY_SELECTOR ,
1414 PROMPT_TEXTAREA_SELECTOR , RESPONSE_CONTAINER_SELECTOR , RESPONSE_TEXT_SELECTOR ,
15- EDIT_MESSAGE_BUTTON_SELECTOR ,USE_URL_CONTEXT_SELECTOR ,UPLOAD_BUTTON_SELECTOR
15+ EDIT_MESSAGE_BUTTON_SELECTOR ,USE_URL_CONTEXT_SELECTOR ,UPLOAD_BUTTON_SELECTOR ,
16+ SET_THINKING_BUDGET_TOGGLE_SELECTOR , THINKING_BUDGET_INPUT_SELECTOR ,
17+ GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR
1618)
1719from config import (
1820 CLICK_TIMEOUT_MS , WAIT_FOR_ELEMENT_TIMEOUT_MS , CLEAR_CHAT_VERIFY_TIMEOUT_MS ,
@@ -62,6 +64,116 @@ async def adjust_parameters(self, request_params: Dict[str, Any], page_params_ca
6264 # 调整URL CONTEXT
6365 await self ._open_url_content (check_client_disconnected )
6466
67+ # 调整“思考预算”开关
68+ await self ._set_thinking_budget_toggle_checked (check_client_disconnected )
69+
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 )
74+
75+ # 调整 Google Search 开关
76+ await self ._adjust_google_search (request_params , check_client_disconnected )
77+
78+ async def _adjust_thinking_budget (self , reasoning_effort : Optional [str ], check_client_disconnected : Callable ):
79+ """根据 reasoning_effort 调整思考预算。"""
80+ self .logger .info (f"[{ self .req_id } )
81+
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 } { 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 :
95+ self .logger .warning (f"[{ self .req_id } { reasoning_effort } )
96+ return
97+
98+ budget_input_locator = self .page .locator (THINKING_BUDGET_INPUT_SELECTOR )
99+
100+ try :
101+ await expect_async (budget_input_locator ).to_be_visible (timeout = 5000 )
102+ await self ._check_disconnect (check_client_disconnected , "思考预算调整 - 输入框可见后" )
103+
104+ self .logger .info (f"[{ self .req_id } { token_budget } )
105+ await budget_input_locator .fill (str (token_budget ), timeout = 5000 )
106+ await self ._check_disconnect (check_client_disconnected , "思考预算调整 - 填充输入框后" )
107+
108+ # 验证
109+ await asyncio .sleep (0.1 )
110+ new_value_str = await budget_input_locator .input_value (timeout = 3000 )
111+ if int (new_value_str ) == token_budget :
112+ self .logger .info (f"[{ self .req_id } { new_value_str } )
113+ else :
114+ self .logger .warning (f"[{ self .req_id } { new_value_str } { token_budget } )
115+
116+ except Exception as e :
117+ self .logger .error (f"[{ self .req_id } { e } )
118+ if isinstance (e , ClientDisconnectedError ):
119+ raise
120+
121+ async def _adjust_google_search (self , request_params : Dict [str , Any ], check_client_disconnected : Callable ):
122+ """根据请求参数中的 'googleSearch' 工具存在与否，双向控制 Google Search 开关。"""
123+ self .logger .info (f"[{ self .req_id } )
124+
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+
133+ toggle_selector = GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR
134+
135+ try :
136+ toggle_locator = self .page .locator (toggle_selector )
137+ await expect_async (toggle_locator ).to_be_visible (timeout = 5000 )
138+ await self ._check_disconnect (check_client_disconnected , "Google Search 开关 - 元素可见后" )
139+
140+ is_checked_str = await toggle_locator .get_attribute ("aria-checked" )
141+ self .logger .info (f"[{ self .req_id } { 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 } )
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 } )
153+ else :
154+ self .logger .warning (f"[{ self .req_id } { new_state } )
155+ else :
156+ self .logger .info (f"[{ self .req_id } )
157+ else :
158+ # 不需要搜索，但开关是打开的，则关闭它
159+ if is_checked_str == "true" :
160+ self .logger .info (f"[{ self .req_id } )
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 } )
167+ else :
168+ self .logger .warning (f"[{ self .req_id } { new_state } )
169+ else :
170+ self .logger .info (f"[{ self .req_id } )
171+
172+ except Exception as e :
173+ self .logger .error (f"[{ self .req_id } { toggle_selector } { e } )
174+ if isinstance (e , ClientDisconnectedError ):
175+ raise
176+
65177 async def _open_url_content (self ,check_client_disconnected : Callable ):
66178 try :
67179 collapse_tools_locator = self .page .locator ('button[aria-label="Expand or collapse tools"]' )
@@ -84,6 +196,53 @@ async def _open_url_content(self,check_client_disconnected: Callable):
84196 except Exception as e :
85197 self .logger .error (f"[{ self .req_id } { e } )
86198
199+ async def _set_thinking_budget_toggle_checked (self , check_client_disconnected : Callable ):
200+ """
201+ 确保 "Thinking Budget" 滑块开关处于选中状态。
202+ """
203+ toggle_selector = SET_THINKING_BUDGET_TOGGLE_SELECTOR
204+ self .logger .info (f"[{ self .req_id } { toggle_selector } )
205+
206+ try :
207+ toggle_locator = self .page .locator (toggle_selector )
208+
209+ # 等待元素在 DOM 中可见
210+ await expect_async (toggle_locator ).to_be_visible (timeout = 5000 )
211+ await self ._check_disconnect (check_client_disconnected , "思考预算开关 - 元素可见后" )
212+
213+ # 1. 检查当前 'aria-checked' 属性
214+ is_checked = await toggle_locator .get_attribute ("aria-checked" )
215+ self .logger .info (f"[{ self .req_id } { is_checked } )
216+
217+ # 2. 如果开关未选中 (aria-checked="false")，则点击它
218+ if is_checked == "false" :
219+ self .logger .info (f"[{ self .req_id } )
220+ 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 } { new_state } )
228+ else :
229+ self .logger .warning (f"[{ self .req_id } { new_state } )
230+ # 如果需要，可以在此处添加错误快照
231+ # await save_error_snapshot(f"thinking_budget_toggle_verify_fail_{self.req_id}")
232+
233+ elif is_checked == "true" :
234+ self .logger .info (f"[{ self .req_id } )
235+ else :
236+ # 处理 'aria-checked' 属性不存在或值无效的情况
237+ self .logger .warning (f"[{ self .req_id } { is_checked } )
238+
239+ except Exception as e :
240+ self .logger .error (f"[{ self .req_id } { toggle_selector } { e } )
241+ # 如果需要，可以在此处添加错误快照
242+ # await save_error_snapshot(f"thinking_budget_toggle_error_{self.req_id}")
243+ # 如果是客户端断开连接的特定错误，则重新抛出
244+ if isinstance (e , ClientDisconnectedError ):
245+ raise
87246 async def _adjust_temperature (self , temperature : float , page_params_cache : dict , params_cache_lock : asyncio .Lock , check_client_disconnected : Callable ):
88247 """调整温度参数。"""
89248 async with params_cache_lock :
0 commit comments