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)
|