223 lines
8.0 KiB
Python
223 lines
8.0 KiB
Python
"""
|
||
豆包网页版驱动（doubao.com/chat/）。
|
||
|
||
生成后不依赖「喇叭瞬时出现/消失」这种不稳定信号。
|
||
而是轮询判断两种输出是否已进入可复制状态：
|
||
1) 右侧写作模式：出现「改用对话直接回答」入口并可点右侧复制按钮
|
||
2) 左侧对话模式：出现可点的消息复制按钮（拿最后一条新增回复）
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
|
||
from .base import BaseEngine
|
||
|
||
_DOUBAO_SPEAKER_PATH_SNIPPET = "M19.8628 9.29346C20.3042"
|
||
|
||
|
||
class DoubaoEngine(BaseEngine):
|
||
async def generate(self, prompt: str) -> str:
|
||
self.page.set_default_timeout(60_000)
|
||
await self.page.goto(
|
||
"https://www.doubao.com/chat/",
|
||
wait_until="domcontentloaded",
|
||
timeout=60_000,
|
||
)
|
||
|
||
editor = self.page.locator('textarea[data-testid="chat_input_input"]').first
|
||
try:
|
||
await editor.wait_for(state="visible", timeout=30_000)
|
||
except Exception:
|
||
return "ERROR:REQUIRE_LOGIN"
|
||
|
||
text = (prompt or "").strip()
|
||
if not text:
|
||
return "ERROR:PROMPT_EMPTY"
|
||
|
||
await editor.click()
|
||
await editor.fill(text)
|
||
await asyncio.sleep(0.2)
|
||
|
||
send = self.page.locator(
|
||
'#flow-end-msg-send, [data-testid="chat_input_send_button"]'
|
||
).first
|
||
try:
|
||
await send.wait_for(state="visible", timeout=15_000)
|
||
await send.click(timeout=10_000)
|
||
except Exception as e:
|
||
return f"ERROR:DOUBAO_SEND_FAILED {e}"
|
||
|
||
print("💡 [llm-manager/doubao] 已发送，等待生成完成后再复制（最长 180s）...")
|
||
|
||
# 发送前记录：用于确认点的是新增复制按钮，且剪贴板确实变化了
|
||
left_copy = self.page.locator('[data-testid="message_action_copy"]')
|
||
right_copy = self.page.locator('[data-testid="container_inner_copy_btn"]')
|
||
left_prev = await left_copy.count()
|
||
right_prev = await right_copy.count()
|
||
clipboard_before = await self._safe_read_clipboard()
|
||
|
||
# 稳定检测：避免生成中“喇叭短暂出现”误判（需要连续多次可见）
|
||
speaker = self.page.locator(
|
||
f'svg path[d*="{_DOUBAO_SPEAKER_PATH_SNIPPET}"]'
|
||
).last
|
||
stable_needed = 3
|
||
stable = 0
|
||
interval_sec = 0.5
|
||
deadline = time.monotonic() + 180.0
|
||
while time.monotonic() < deadline:
|
||
try:
|
||
if await speaker.is_visible():
|
||
stable += 1
|
||
else:
|
||
stable = 0
|
||
except Exception:
|
||
stable = 0
|
||
if stable >= stable_needed:
|
||
break
|
||
await asyncio.sleep(interval_sec)
|
||
else:
|
||
return "ERROR:DOUBAO_WAIT_COMPLETE_TIMEOUT 180 秒内未检测到“喇叭稳定回显”。"
|
||
|
||
# 生成完成后：依据当前页面状态决定点右侧还是左侧复制
|
||
write_mode = await self._is_write_mode()
|
||
if write_mode:
|
||
if await self._wait_count_gt(right_copy, right_prev, timeout_sec=15.0):
|
||
copy_btn = right_copy.last
|
||
elif await self._wait_count_gt(left_copy, left_prev, timeout_sec=15.0):
|
||
copy_btn = left_copy.last
|
||
else:
|
||
return "ERROR:DOUBAO_COPY_BUTTON_NOT_READY 右侧/左侧复制按钮均未出现新增。"
|
||
else:
|
||
if await self._wait_count_gt(left_copy, left_prev, timeout_sec=15.0):
|
||
copy_btn = left_copy.last
|
||
elif await self._wait_count_gt(right_copy, right_prev, timeout_sec=5.0):
|
||
copy_btn = right_copy.last
|
||
else:
|
||
return "ERROR:DOUBAO_COPY_BUTTON_NOT_READY 左侧复制按钮未出现新增。"
|
||
|
||
try:
|
||
await copy_btn.scroll_into_view_if_needed()
|
||
await copy_btn.click(timeout=15_000, force=True)
|
||
except Exception as e:
|
||
return f"ERROR:DOUBAO_COPY_CLICK_FAILED {e}"
|
||
|
||
out = await self._wait_clipboard_changed(clipboard_before, timeout_sec=20.0)
|
||
if out is None:
|
||
return "ERROR:DOUBAO_CLIPBOARD_NOT_UPDATED 剪贴板未在规定时间内更新。"
|
||
if self._is_invalid_clipboard(out, clipboard_before=clipboard_before):
|
||
return f"ERROR:DOUBAO_CLIPBOARD_INVALID {out[:200]}"
|
||
|
||
await asyncio.sleep(5)
|
||
return out
|
||
|
||
async def _is_write_mode(self) -> bool:
|
||
marker = self.page.locator('div.bottom-entry-GSWErB:has-text("改用对话直接回答")').first
|
||
try:
|
||
return await marker.is_visible()
|
||
except Exception:
|
||
return False
|
||
|
||
async def _wait_count_gt(self, locator, prev_count: int, timeout_sec: float) -> bool:
|
||
deadline = time.monotonic() + timeout_sec
|
||
while time.monotonic() < deadline:
|
||
try:
|
||
if await locator.count() > prev_count:
|
||
return True
|
||
except Exception:
|
||
pass
|
||
await asyncio.sleep(0.3)
|
||
return False
|
||
|
||
async def _wait_clipboard_changed(
|
||
self, clipboard_before: str | None, timeout_sec: float
|
||
) -> str | None:
|
||
deadline = time.monotonic() + timeout_sec
|
||
last = clipboard_before
|
||
while time.monotonic() < deadline:
|
||
cur = await self._safe_read_clipboard()
|
||
if cur and cur != last and not self._is_invalid_clipboard(
|
||
cur, clipboard_before=clipboard_before
|
||
):
|
||
return cur.strip()
|
||
last = cur
|
||
await asyncio.sleep(0.4)
|
||
return None
|
||
|
||
async def _try_copy_right(self, right_copy) -> str | None:
|
||
# 右侧复制按钮
|
||
try:
|
||
btn = right_copy.first
|
||
await btn.wait_for(state="visible", timeout=2000)
|
||
await btn.scroll_into_view_if_needed()
|
||
await btn.click(timeout=5000, force=True)
|
||
except Exception:
|
||
return None
|
||
|
||
await asyncio.sleep(0.4)
|
||
out = await self._safe_read_clipboard()
|
||
if self._is_invalid_clipboard(out):
|
||
return None
|
||
if out == (await self._safe_read_clipboard()):
|
||
pass
|
||
return (out or "").strip()
|
||
|
||
async def _try_copy_left(self, left_copy, prev_count: int) -> str | None:
|
||
n = await left_copy.count()
|
||
if n <= prev_count:
|
||
return None
|
||
|
||
# 从新增区间末尾往前找第一个可点击的复制按钮
|
||
for idx in range(n - 1, prev_count - 1, -1):
|
||
btn = left_copy.nth(idx)
|
||
try:
|
||
if not await btn.is_visible():
|
||
continue
|
||
await btn.scroll_into_view_if_needed()
|
||
await btn.click(timeout=5000, force=True)
|
||
except Exception:
|
||
continue
|
||
|
||
await asyncio.sleep(0.4)
|
||
out = await self._safe_read_clipboard()
|
||
if self._is_invalid_clipboard(out):
|
||
continue
|
||
return (out or "").strip()
|
||
return None
|
||
|
||
def _is_invalid_clipboard(
|
||
self, text: str | None, clipboard_before: str | None = None
|
||
) -> bool:
|
||
if not text:
|
||
return True
|
||
t = str(text).strip()
|
||
if not t:
|
||
return True
|
||
if clipboard_before is not None and t == str(clipboard_before).strip():
|
||
return True
|
||
if t.startswith("ERROR:"):
|
||
return True
|
||
# 防止误点拿到 UI 文案/占位文本
|
||
banned_substrings = (
|
||
"改用对话直接回答",
|
||
"新对话",
|
||
"正在生成",
|
||
"复制",
|
||
"下载",
|
||
"快捷",
|
||
"发消息",
|
||
)
|
||
if any(s in t for s in banned_substrings):
|
||
return True
|
||
if "data-testid" in t or "<button" in t:
|
||
return True
|
||
# 经验阈值：正常正文不会太短
|
||
if len(t) < 30:
|
||
return True
|
||
return False
|
||
|
||
async def _safe_read_clipboard(self) -> str | None:
|
||
try:
|
||
return await self.page.evaluate("navigator.clipboard.readText()")
|
||
except Exception:
|
||
return None
|