397 lines
13 KiB
Python
397 lines
13 KiB
Python
|
import asyncio
|
|||
|
|
import random
|
|||
|
|
import time
|
|||
|
|
from typing import AsyncGenerator, List, Optional, Tuple
|
|||
|
|
|
|||
|
|
from loguru import logger
|
|||
|
|
from playwright.async_api import Locator, Page, TimeoutError as PWTimeout
|
|||
|
|
|
|||
|
|
from .config import settings
|
|||
|
|
from .models import ChatMessage
|
|||
|
|
|
|||
|
|
FIND_INPUT_JS = """
|
|||
|
|
() => {
|
|||
|
|
const selectors = [
|
|||
|
|
'.tiptap.ProseMirror[contenteditable="true"]',
|
|||
|
|
'div[contenteditable="true"]',
|
|||
|
|
'textarea[placeholder*="发"]',
|
|||
|
|
'textarea[placeholder*="输入"]',
|
|||
|
|
'textarea[placeholder*="问"]',
|
|||
|
|
'textarea',
|
|||
|
|
'[role="textbox"]',
|
|||
|
|
];
|
|||
|
|
for (const sel of selectors) {
|
|||
|
|
const el = document.querySelector(sel);
|
|||
|
|
if (el && el.offsetParent !== null) {
|
|||
|
|
const rect = el.getBoundingClientRect();
|
|||
|
|
if (rect.width > 0 && rect.height > 0) return sel;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return null;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
FIND_SEND_BTN_JS = """
|
|||
|
|
() => {
|
|||
|
|
const selectors = [
|
|||
|
|
'button[aria-label*="发送"]',
|
|||
|
|
'button[aria-label*="Send"]',
|
|||
|
|
'button[type="submit"]',
|
|||
|
|
'button[class*="send"]',
|
|||
|
|
'button[class*="submit"]',
|
|||
|
|
'[data-testid*="send"]',
|
|||
|
|
'button[class*="Send"]',
|
|||
|
|
];
|
|||
|
|
for (const sel of selectors) {
|
|||
|
|
const el = document.querySelector(sel);
|
|||
|
|
if (el && el.offsetParent !== null) return sel;
|
|||
|
|
}
|
|||
|
|
return null;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
GET_RESPONSE_TEXT_JS = """
|
|||
|
|
() => {
|
|||
|
|
const messageList = document.querySelector('[class*="message-list"]');
|
|||
|
|
if (messageList) {
|
|||
|
|
const msgs = Array.from(messageList.children).filter(el => {
|
|||
|
|
const cls = el.className || '';
|
|||
|
|
return el.offsetParent !== null
|
|||
|
|
&& !cls.includes('action-bar')
|
|||
|
|
&& !cls.includes('suggest-')
|
|||
|
|
&& el.innerText.trim();
|
|||
|
|
});
|
|||
|
|
if (msgs.length > 0) {
|
|||
|
|
return msgs[msgs.length - 1].innerText.trim();
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return null;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
IS_GENERATING_JS = """
|
|||
|
|
() => {
|
|||
|
|
const stopSelectors = [
|
|||
|
|
'[class*="stop-generate"]',
|
|||
|
|
'[class*="stop_generate"]',
|
|||
|
|
'[aria-label*="停止"]',
|
|||
|
|
'[aria-label*="Stop"]',
|
|||
|
|
'[data-testid*="stop"]',
|
|||
|
|
'button[disabled][class*="send"]',
|
|||
|
|
];
|
|||
|
|
for (const sel of stopSelectors) {
|
|||
|
|
const el = document.querySelector(sel);
|
|||
|
|
if (el && el.offsetParent !== null) return true;
|
|||
|
|
}
|
|||
|
|
const spinnerSelectors = [
|
|||
|
|
'[class*="typing-indicator"]',
|
|||
|
|
'[class*="thinking"]',
|
|||
|
|
'[class*="loading-dot"]',
|
|||
|
|
'svg[class*="loading"]',
|
|||
|
|
'[class*="generating"]',
|
|||
|
|
];
|
|||
|
|
for (const sel of spinnerSelectors) {
|
|||
|
|
const el = document.querySelector(sel);
|
|||
|
|
if (el && el.offsetParent !== null) return true;
|
|||
|
|
}
|
|||
|
|
return false;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
NEW_CHAT_BTN_JS = """
|
|||
|
|
() => {
|
|||
|
|
const selectors = [
|
|||
|
|
'button[aria-label*="新对话"]',
|
|||
|
|
'button[aria-label*="New"]',
|
|||
|
|
'a[aria-label*="新对话"]',
|
|||
|
|
'a[aria-label*="New"]',
|
|||
|
|
'[class*="new-chat"]',
|
|||
|
|
'[class*="new_chat"]',
|
|||
|
|
'[data-testid*="new-chat"]',
|
|||
|
|
'[data-testid*="new_chat"]',
|
|||
|
|
'a[href*="/chat"][class*="new"]',
|
|||
|
|
];
|
|||
|
|
for (const sel of selectors) {
|
|||
|
|
const el = document.querySelector(sel);
|
|||
|
|
if (el && el.offsetParent !== null) return sel;
|
|||
|
|
}
|
|||
|
|
return null;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
INSPECT_DOM_JS = """
|
|||
|
|
() => {
|
|||
|
|
const result = {
|
|||
|
|
url: window.location.href,
|
|||
|
|
title: document.title,
|
|||
|
|
textareas: [],
|
|||
|
|
contentEditables: [],
|
|||
|
|
buttons: [],
|
|||
|
|
messageContainers: [],
|
|||
|
|
};
|
|||
|
|
document.querySelectorAll('textarea').forEach(el => {
|
|||
|
|
result.textareas.push({
|
|||
|
|
placeholder: el.placeholder,
|
|||
|
|
className: el.className.slice(0, 80),
|
|||
|
|
visible: el.offsetParent !== null,
|
|||
|
|
});
|
|||
|
|
});
|
|||
|
|
document.querySelectorAll('[contenteditable="true"]').forEach(el => {
|
|||
|
|
result.contentEditables.push({
|
|||
|
|
className: el.className.slice(0, 80),
|
|||
|
|
visible: el.offsetParent !== null,
|
|||
|
|
});
|
|||
|
|
});
|
|||
|
|
document.querySelectorAll('button').forEach(el => {
|
|||
|
|
if (el.offsetParent !== null) {
|
|||
|
|
result.buttons.push({
|
|||
|
|
text: el.innerText.slice(0, 40),
|
|||
|
|
ariaLabel: el.getAttribute('aria-label'),
|
|||
|
|
className: el.className.slice(0, 80),
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
});
|
|||
|
|
const msgPatterns = ['message', 'chat-item', 'conversation', 'receive', 'bot', 'ai-msg'];
|
|||
|
|
msgPatterns.forEach(pattern => {
|
|||
|
|
document.querySelectorAll(`[class*="${pattern}"]`).forEach(el => {
|
|||
|
|
if (el.offsetParent !== null && result.messageContainers.length < 20) {
|
|||
|
|
result.messageContainers.push({
|
|||
|
|
pattern: pattern,
|
|||
|
|
className: el.className.slice(0, 80),
|
|||
|
|
text: el.innerText.slice(0, 100),
|
|||
|
|
childCount: el.children.length,
|
|||
|
|
});
|
|||
|
|
}
|
|||
|
|
});
|
|||
|
|
});
|
|||
|
|
return result;
|
|||
|
|
}
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DoubaoChat:
|
|||
|
|
def __init__(self, page: Page):
|
|||
|
|
self.page = page
|
|||
|
|
|
|||
|
|
async def start_new_chat(self) -> bool:
|
|||
|
|
"""Try to start a new chat conversation."""
|
|||
|
|
try:
|
|||
|
|
selector = await self.page.evaluate(NEW_CHAT_BTN_JS)
|
|||
|
|
if selector:
|
|||
|
|
btn = self.page.locator(selector).first
|
|||
|
|
await btn.click()
|
|||
|
|
await self.page.wait_for_timeout(2000)
|
|||
|
|
logger.info("Started new chat via button")
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.debug(f"New chat button not found: {e}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
await self.page.goto(settings.doubao_url, wait_until="domcontentloaded")
|
|||
|
|
await self.page.wait_for_timeout(3000)
|
|||
|
|
logger.info("Navigated to fresh chat URL")
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"Failed to navigate to chat URL: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def _find_input(self) -> Optional[Tuple[Locator, str]]:
|
|||
|
|
"""Find the input element, returns (locator, selector)."""
|
|||
|
|
selector = await self.page.evaluate(FIND_INPUT_JS)
|
|||
|
|
if not selector:
|
|||
|
|
return None
|
|||
|
|
loc = self.page.locator(selector).first
|
|||
|
|
try:
|
|||
|
|
await loc.wait_for(state="visible", timeout=5000)
|
|||
|
|
return loc, selector
|
|||
|
|
except PWTimeout:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def _find_send_button(self) -> Optional[Locator]:
|
|||
|
|
"""Find the send button."""
|
|||
|
|
selector = await self.page.evaluate(FIND_SEND_BTN_JS)
|
|||
|
|
if not selector:
|
|||
|
|
return None
|
|||
|
|
loc = self.page.locator(selector).first
|
|||
|
|
try:
|
|||
|
|
await loc.wait_for(state="visible", timeout=3000)
|
|||
|
|
return loc
|
|||
|
|
except PWTimeout:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def _type_message(self, input_loc: Locator, text: str):
|
|||
|
|
"""Type a message with human-like delays."""
|
|||
|
|
await input_loc.click()
|
|||
|
|
await asyncio.sleep(0.2)
|
|||
|
|
|
|||
|
|
if len(text) > 1000:
|
|||
|
|
await self.page.keyboard.insert_text(text)
|
|||
|
|
await asyncio.sleep(0.3)
|
|||
|
|
else:
|
|||
|
|
delay = random.randint(settings.typing_delay_min, settings.typing_delay_max)
|
|||
|
|
await self.page.keyboard.type(text, delay=delay)
|
|||
|
|
|
|||
|
|
await asyncio.sleep(settings.inter_message_delay)
|
|||
|
|
|
|||
|
|
async def _send_via_enter(self, input_loc: Locator) -> bool:
|
|||
|
|
"""Try sending by pressing Enter."""
|
|||
|
|
await input_loc.press("Enter")
|
|||
|
|
await asyncio.sleep(0.5)
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def _send_via_button(self) -> bool:
|
|||
|
|
"""Try sending by clicking the send button."""
|
|||
|
|
btn = await self._find_send_button()
|
|||
|
|
if btn:
|
|||
|
|
await btn.click()
|
|||
|
|
await asyncio.sleep(0.5)
|
|||
|
|
return True
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def _is_input_empty(self, input_loc: Locator, selector: str) -> bool:
|
|||
|
|
"""Check if the input element is empty (message was sent)."""
|
|||
|
|
try:
|
|||
|
|
if "textarea" in selector:
|
|||
|
|
val = await input_loc.input_value()
|
|||
|
|
else:
|
|||
|
|
val = await input_loc.inner_text()
|
|||
|
|
return not val.strip()
|
|||
|
|
except Exception:
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def send_message(self, text: str) -> str:
|
|||
|
|
"""Send a message to the Doubao chat.
|
|||
|
|
Returns the previous response text (for change detection).
|
|||
|
|
"""
|
|||
|
|
result = await self._find_input()
|
|||
|
|
if not result:
|
|||
|
|
logger.error("Input element not found on page")
|
|||
|
|
raise RuntimeError("无法找到输入框，请通过VNC检查页面状态")
|
|||
|
|
|
|||
|
|
input_loc, selector = result
|
|||
|
|
|
|||
|
|
old_response = await self.page.evaluate(GET_RESPONSE_TEXT_JS) or ""
|
|||
|
|
|
|||
|
|
await self._type_message(input_loc, text)
|
|||
|
|
|
|||
|
|
sent = await self._send_via_enter(input_loc)
|
|||
|
|
if sent:
|
|||
|
|
await asyncio.sleep(0.8)
|
|||
|
|
if not await self._is_input_empty(input_loc, selector):
|
|||
|
|
logger.info("Enter didn't send, trying send button")
|
|||
|
|
sent = await self._send_via_button()
|
|||
|
|
|
|||
|
|
if not sent:
|
|||
|
|
logger.error("Failed to send message")
|
|||
|
|
raise RuntimeError("消息发送失败，请通过VNC检查页面状态")
|
|||
|
|
|
|||
|
|
return old_response
|
|||
|
|
|
|||
|
|
async def _wait_for_response_start(
|
|||
|
|
self, old_text: str, timeout: int = 30
|
|||
|
|
) -> Optional[str]:
|
|||
|
|
"""Wait for the response to start appearing.
|
|||
|
|
|
|||
|
|
Uses two-phase detection:
|
|||
|
|
Phase 1: wait for user message to appear (text changes from old_text)
|
|||
|
|
Phase 2: wait for AI response to start (text changes from baseline)
|
|||
|
|
|
|||
|
|
Returns baseline text (after user message, before AI response),
|
|||
|
|
or None on timeout.
|
|||
|
|
"""
|
|||
|
|
start = time.time()
|
|||
|
|
phase = 1
|
|||
|
|
baseline = None
|
|||
|
|
|
|||
|
|
while time.time() - start < timeout:
|
|||
|
|
current = await self.page.evaluate(GET_RESPONSE_TEXT_JS) or ""
|
|||
|
|
|
|||
|
|
if phase == 1:
|
|||
|
|
if current and current != old_text:
|
|||
|
|
baseline = current
|
|||
|
|
phase = 2
|
|||
|
|
await asyncio.sleep(0.5)
|
|||
|
|
elif phase == 2:
|
|||
|
|
if current != baseline:
|
|||
|
|
logger.debug("Response started")
|
|||
|
|
return baseline
|
|||
|
|
|
|||
|
|
await asyncio.sleep(0.3)
|
|||
|
|
|
|||
|
|
if baseline is not None:
|
|||
|
|
logger.warning("Response may have started during baseline capture")
|
|||
|
|
return baseline
|
|||
|
|
|
|||
|
|
logger.warning("Response did not start within timeout")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def stream_response(self, old_text: str = "") -> AsyncGenerator[str, None]:
|
|||
|
|
"""Yield response text chunks as they appear.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
old_text: Previous response text from before the message was sent.
|
|||
|
|
Used to detect when the new response starts.
|
|||
|
|
"""
|
|||
|
|
baseline = await self._wait_for_response_start(old_text)
|
|||
|
|
if baseline is None:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
last_text = baseline
|
|||
|
|
stable_count = 0
|
|||
|
|
STABLE_THRESHOLD = 6
|
|||
|
|
POLL_INTERVAL = 0.15
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
current_text = await self.page.evaluate(GET_RESPONSE_TEXT_JS) or ""
|
|||
|
|
|
|||
|
|
if current_text and len(current_text) > len(last_text):
|
|||
|
|
new_chunk = current_text[len(last_text):]
|
|||
|
|
last_text = current_text
|
|||
|
|
stable_count = 0
|
|||
|
|
if new_chunk.strip():
|
|||
|
|
yield new_chunk
|
|||
|
|
|
|||
|
|
elif current_text and current_text == last_text:
|
|||
|
|
stable_count += 1
|
|||
|
|
if stable_count >= STABLE_THRESHOLD:
|
|||
|
|
is_gen = await self.page.evaluate(IS_GENERATING_JS)
|
|||
|
|
if not is_gen:
|
|||
|
|
logger.debug("Response complete (stable + not generating)")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
elapsed = time.time() - start_time
|
|||
|
|
if elapsed > settings.response_timeout:
|
|||
|
|
logger.warning(f"Response timeout after {elapsed:.1f}s")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
await asyncio.sleep(POLL_INTERVAL)
|
|||
|
|
|
|||
|
|
async def get_response_full(self, old_text: str = "") -> str:
|
|||
|
|
"""Get the complete response text (non-streaming)."""
|
|||
|
|
chunks = []
|
|||
|
|
async for chunk in self.stream_response(old_text):
|
|||
|
|
chunks.append(chunk)
|
|||
|
|
return "".join(chunks)
|
|||
|
|
|
|||
|
|
async def inspect(self) -> dict:
|
|||
|
|
"""Return DOM structure info for debugging."""
|
|||
|
|
return await self.page.evaluate(INSPECT_DOM_JS)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def messages_to_prompt(messages: List[ChatMessage]) -> str:
|
|||
|
|
"""Convert OpenAI messages array to a single text prompt."""
|
|||
|
|
if len(messages) == 1 and messages[0].role == "user":
|
|||
|
|
return messages[0].content
|
|||
|
|
|
|||
|
|
parts: List[str] = []
|
|||
|
|
for msg in messages:
|
|||
|
|
if msg.role == "system":
|
|||
|
|
parts.append(f"【系统设定】\n{msg.content}")
|
|||
|
|
elif msg.role == "user":
|
|||
|
|
parts.append(f"【用户】\n{msg.content}")
|
|||
|
|
elif msg.role == "assistant":
|
|||
|
|
parts.append(f"【AI回复】\n{msg.content}")
|
|||
|
|
|
|||
|
|
return "\n\n---\n\n".join(parts)
|