acquire-lead/collectors/douyin/ai.py

"""步骤：AI抖音总结。"""

from __future__ import annotations

import re
import time

from playwright.sync_api import Page

from collectors.douyin.common import AI_SUMMARY_PROMPT, human_pause
from collectors.douyin.panel import _click_side_panel_tab

def _get_ai_frame(page: Page):
    """获取 AI抖音 iframe（so-landing / search_ai_mobile）。"""
    deadline = time.time() + 15
    while time.time() < deadline:
        for frame in page.frames:
            url = frame.url or ""
            if "search_ai" in url or "so-landing.douyin.com" in url:
                return frame
        try:
            page.wait_for_selector(
                'iframe[src*="search_ai_mobile"], iframe.D99qPWMG',
                timeout=2_000,
            )
        except Exception:
            human_pause(0.5, 0.9)
    return None


def _parse_ai_summary_text(body_text: str, prompt: str) -> str:
    """从 AI iframe 文本中提取总结正文。"""
    text = (body_text or "").strip()
    if not text:
        return ""

    # 优先取「最后一次提问」之后的回复
    if prompt and prompt in text:
        reply = text.rsplit(prompt, 1)[-1].strip()
    else:
        reply = text

    # 若存在「以上为历史记录」，取其后内容
    if "以上为历史记录" in reply:
        reply = reply.split("以上为历史记录", 1)[-1].strip()
        if prompt and prompt in reply:
            reply = reply.split(prompt, 1)[-1].strip()

    skip_exact = {
        "AI 总结",
        "高光片段",
        "换一换",
        "下载抖音精选",
        "以上为历史记录",
    }
    skip_prefix = (
        "以上结果",
        "看到视频",
        "分享一些",
        "有哪些",
        "如何",
        "推荐一些",
        "核心原理",
        "无效输入",
        "i+1",
        "自学可行性",
        "视频提到",
    )
    lines = [ln.strip() for ln in reply.splitlines() if ln.strip()]
    for line in lines:
        if line in skip_exact:
            continue
        if line.startswith(skip_prefix):
            continue
        if re.match(r"^\d{1,2}:\d{2}", line):
            continue
        if line == prompt:
            continue
        if 8 <= len(line) <= 200:
            return line[:512]

    # 回退：页面上「AI 总结」后的首段
    if "AI 总结" in text:
        after = text.split("AI 总结", 1)[-1]
        for line in after.splitlines():
            line = line.strip()
            if 8 <= len(line) <= 200 and line not in skip_exact:
                return line[:512]
    return ""


def collect_ai_summary(page: Page, timeout_ms: int = 60_000) -> str:
    """
    前提：已点击评论图标打开右侧栏。
    切换到「AI抖音」→ iframe 输入发送 → 返回总结文本。
    """
    print("[AI抖音] 切换到 AI抖音 tab...")
    if not _click_side_panel_tab(page, "AI抖音"):
        raise RuntimeError("无法切换到「AI抖音」")

    frame = _get_ai_frame(page)
    if frame is None:
        raise RuntimeError("未找到 AI抖音 iframe")
    try:
        frame.wait_for_load_state("domcontentloaded")
    except Exception:
        pass

    frame.wait_for_selector('[contenteditable="true"], .input_blVmyq', timeout=20_000)
    human_pause(0.8, 1.2)
    before_text = frame.evaluate("() => document.body.innerText || ''") or ""

    box = frame.locator('[contenteditable="true"], .input_blVmyq').first
    box.click(timeout=8_000)
    page.keyboard.press("Control+A")
    page.keyboard.press("Backspace")
    page.keyboard.type(AI_SUMMARY_PROMPT, delay=30)
    human_pause(0.6, 1.0)

    typed = (
        frame.evaluate(
            """() => {
                const el = document.querySelector('[contenteditable="true"], .input_blVmyq');
                return el ? (el.innerText || el.textContent || '').trim() : '';
            }"""
        )
        or ""
    )
    if not typed:
        raise RuntimeError("AI抖音输入框未能写入总结指令")
    print(f"[AI抖音] 已输入: {typed[:40]}")

    # 等发送按钮激活再点
    sent = False
    for _ in range(10):
        sent = bool(
            frame.evaluate(
                """() => {
                    const btn = document.querySelector('.searchButtonActive_viYBkl')
                        || document.querySelector('.searchButton_IMT5UH');
                    if (!btn) return false;
                    const r = btn.getBoundingClientRect();
                    if (r.width <= 0 || r.height <= 0) return false;
                    const msg = btn.getAttribute('data-msg') || '';
                    if (/上传|截图|图片/.test(msg)) return false;
                    btn.click();
                    return true;
                }"""
            )
        )
        if sent:
            break
        time.sleep(0.35)

    if not sent:
        page.keyboard.press("Enter")
        print("[AI抖音] 未定位发送按钮，已回车发送")
    else:
        print("[AI抖音] 已发送，等待回复...")

    deadline = time.time() + timeout_ms / 1000
    summary = ""
    while time.time() < deadline:
        human_pause(1.2, 1.8)
        body = frame.evaluate("() => document.body.innerText || ''") or ""
        candidate = _parse_ai_summary_text(body, AI_SUMMARY_PROMPT)
        if not candidate:
            continue
        if "以上为历史记录" in body or "以上结果由问问AI" in body:
            summary = candidate
            break
        if candidate not in (before_text or ""):
            summary = candidate
            break

    if not summary:
        body = frame.evaluate("() => document.body.innerText || ''") or ""
        summary = _parse_ai_summary_text(body, AI_SUMMARY_PROMPT)
    if not summary:
        raise RuntimeError("AI抖音未返回总结内容")

    summary = summary.strip()[:512]
    print(f"[AI抖音] 总结: {summary}")
    return summary