holy-python/crawler_client.py

import re
import secrets
import time
from typing import Any

import requests
import json
import hashlib

from config import Config


class CrawlerError(RuntimeError):
    pass


class NoAiCardError(CrawlerError):
    """搜索请求成功，但当前结果中没有问一问 AI 卡片。"""


CRAWLER_SESSION = requests.Session()
# 爬虫服务为明确配置的直连地址，不应继承 macOS 系统代理。
# 否则 requests 会经由 127.0.0.1:7890，约 30 秒后被代理返回 502。
CRAWLER_SESSION.trust_env = False


def normalize_collection_keywords(keywords: list[str]) -> list[str]:
    """清理影响小红书问一问卡片匹配的末尾标点，并保持关键词顺序。"""

    normalized: list[str] = []
    for keyword in keywords:
        value = re.sub(r"[?？!！。．.]+$", "", str(keyword).strip()).strip()
        if value and value not in normalized:
            normalized.append(value)
    if not normalized:
        raise CrawlerError("关键词不能为空")
    return normalized


def compact_search_keyword(keyword: str) -> str:
    """生成问一问卡片匹配的无空白回退词。

    小红书会把 ``ai 生成 ppt 哪个好用`` 和 ``ai生成ppt哪个好用`` 当成
    不同搜索词，前者可能不返回 AI 卡片。仅在原词明确无卡片时使用该
    回退，不改变任务对外展示的原始关键词。
    """

    return re.sub(r"\s+", "", keyword)


def fallback_search_keywords(keyword: str) -> list[str]:
    """Generate conservative query variants when the original query misses the AI card."""

    variants: list[str] = []

    def add(value: str) -> None:
        value = re.sub(r"\s+", " ", value).strip()
        if value and value != keyword and value not in variants:
            variants.append(value)

    add(compact_search_keyword(keyword))

    # “aippt自动生成哪个好”在搜索端可能无法正确分词，而等价的自然分词
    # “ai 生成 ppt 哪个好用”可以命中相同意图下的问一问 AI 卡片。
    spaced = re.sub(
        r"^ai\s*ppt\s*自动生成",
        "ai 生成 ppt ",
        keyword,
        flags=re.IGNORECASE,
    )
    spaced = re.sub(r"哪个好$", "哪个好用", spaced)
    add(spaced)
    return variants


def create_search_context_id() -> str:
    """生成小红书搜索链路使用的 36 进制上下文 ID。

    与 Web 端 createSearchId 的算法保持一致：
    ``(Date.now() << 64) + random(1..2147483646)``。
    """

    value = (int(time.time() * 1000) << 64) + secrets.randbelow(2147483646) + 1
    alphabet = "0123456789abcdefghijklmnopqrstuvwxyz"
    encoded = ""
    while value:
        value, remainder = divmod(value, 36)
        encoded = alphabet[remainder] + encoded
    return encoded or "0"


def _post(
    path: str,
    body: dict[str, Any],
    *,
    retry_transient: bool = True,
) -> dict[str, Any]:
    url = f"{Config.CRAWLER_BASE_URL}{path}"
    retry_times = max(1, Config.CRAWLER_RETRY_TIMES) if retry_transient else 1
    response: requests.Response | None = None

    for attempt in range(1, retry_times + 1):
        try:
            response = CRAWLER_SESSION.post(
                url,
                json=body,
                timeout=Config.CRAWLER_TIMEOUT_SECONDS,
            )
            # 网关临时异常时自动重试，避免直接将业务任务标记为失败。
            if response.status_code in {502, 503, 504} and attempt < retry_times:
                time.sleep(Config.CRAWLER_RETRY_DELAY_SECONDS * attempt)
                continue
            response.raise_for_status()
            payload = response.json()
            break
        except (requests.ConnectionError, requests.Timeout) as exc:
            if attempt < retry_times:
                time.sleep(Config.CRAWLER_RETRY_DELAY_SECONDS * attempt)
                continue
            raise CrawlerError(f"爬虫服务请求失败（已重试 {retry_times} 次）：{exc}") from exc
        except requests.RequestException as exc:
            detail = response.text[:500] if response is not None else ""
            message = f"爬虫服务请求失败：{exc}"
            if detail:
                message += f"，响应：{detail}"
            raise CrawlerError(message) from exc
        except ValueError as exc:
            raise CrawlerError("爬虫服务返回的不是有效 JSON") from exc
    else:
        raise CrawlerError(f"爬虫服务请求失败（已重试 {retry_times} 次）")

    if payload.get("code") != 200 or not payload.get("success", False):
        message = str(payload.get("msg") or "爬虫服务返回失败")
        no_ai_card_signals = (
            "没有返回问一问 AI 卡片",
            "NoneType' object has no attribute 'get'",
            "complex/detail 未返回 AI 数据",
        )
        if any(signal in message for signal in no_ai_card_signals):
            raise NoAiCardError(message)
        raise CrawlerError(message)
    data = payload.get("data")
    if not isinstance(data, dict):
        raise CrawlerError("爬虫服务响应缺少 data")
    return data


def start_collection(mode: str, keywords: list[str]) -> dict[str, Any]:
    """深度任务走 sources，常规任务走 overview。"""

    collection_keywords = normalize_collection_keywords(keywords)
    path = "/api/v1/xhs/wen/sources" if mode == "deep" else "/api/v1/xhs/wen/overview"
    # 常规任务允许批量关键词；数组由远程 overview 创建异步批量任务。
    keyword: str | list[str] = collection_keywords if len(collection_keywords) > 1 else collection_keywords[0]
    # 轻度 overview 在首次计算超过网关时限后，远程通常已经写好缓存，
    # 允许受控重试以取得结果；深度 sources 会启动长任务，禁止盲目重试，
    # 避免远程请求虽断开但后台继续执行时产生重复采集任务。
    # 采集服务的 search/notes 需要搜索上下文，但业务调用方只应传 keyword。
    # 在 Holy蟹内部为每个新任务生成一次，并由同一个远程采集任务全程复用。
    def submit(value: str | list[str]) -> dict[str, Any]:
        return _post(
            path,
            {
                "keyword": value,
                "search_id": create_search_context_id(),
                "session_id": create_search_context_id(),
            },
            retry_transient=(mode != "deep"),
        )

    try:
        return submit(keyword)
    except NoAiCardError:
        # 搜索端分词差异可能让真实存在的问一问卡片暂时无法命中。
        # 依次尝试保守的等价写法；全部失败后才交给上层按空结果完成。
        if isinstance(keyword, str):
            for fallback_keyword in fallback_search_keywords(keyword):
                try:
                    result = submit(fallback_keyword)
                except NoAiCardError:
                    continue
                result["requested_keyword"] = keyword
                result["collection_keyword"] = fallback_keyword
                return result
        raise


def query_collection(crawler_task_id: str) -> dict[str, Any]:
    return _post("/api/v1/xhs/wen/tasks/query", {"task_id": crawler_task_id})


def recover_source_note_details(keyword: str) -> dict[str, Any]:
    """从采集器已落库的 complex/source 记录恢复可分析的来源 notes。"""

    if not all((
        Config.CRAWLER_DATABASE_HOST,
        Config.CRAWLER_DATABASE_USER,
        Config.CRAWLER_DATABASE_NAME,
    )):
        return {}
    import pymysql

    db = pymysql.connect(
        host=Config.CRAWLER_DATABASE_HOST,
        port=Config.CRAWLER_DATABASE_PORT,
        user=Config.CRAWLER_DATABASE_USER,
        password=Config.CRAWLER_DATABASE_PASSWORD,
        database=Config.CRAWLER_DATABASE_NAME,
        charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor,
    )
    try:
        with db.cursor() as cursor:
            cursor.execute(
                "SELECT searchid,card_id FROM xhs_wen_ai_answers "
                "WHERE keyword=%s ORDER BY updated_at DESC,id DESC LIMIT 1",
                (keyword,),
            )
            linkage = cursor.fetchone()
            if not linkage:
                return {}
            cursor.execute(
                "SELECT n.raw_json,n.tag_name,t.title AS product_name "
                "FROM xhs_wen_tag_notes n LEFT JOIN xhs_wen_ai_tags t "
                "ON t.keyword=n.keyword AND t.searchid=n.searchid "
                "AND t.card_id=n.card_id AND t.entity_id=n.entity_id "
                "AND t.tag_name=n.tag_name "
                "WHERE n.keyword=%s AND n.searchid=%s AND n.card_id=%s "
                "ORDER BY n.id",
                (keyword, linkage["searchid"], linkage["card_id"]),
            )
            rows = cursor.fetchall()
    finally:
        db.close()

    grouped: dict[str, dict[str, list[dict[str, Any]]]] = {}
    for row in rows:
        try:
            source = json.loads(row.get("raw_json") or "{}")
        except (TypeError, ValueError):
            continue
        note = source.get("note_info") if isinstance(source.get("note_info"), dict) else {}
        if not note:
            continue
        note = dict(note)
        note.setdefault("type", source.get("type") or note.get("type"))
        note.setdefault(
            "wechat_share_desc",
            source.get("wechat_share_desc")
            or note.get("paragraph_text")
            or note.get("relevant_text"),
        )
        product = str(row.get("product_name") or "")
        tag = str(row.get("tag_name") or "")
        grouped.setdefault(product, {}).setdefault(tag, []).append(note)

    products = [
        {
            "product_name": product,
            "content_tags": [
                {"content_tag": tag, "count": len(notes), "notes": notes}
                for tag, notes in tags.items()
            ],
        }
        for product, tags in grouped.items()
    ]
    count = sum(
        len(tag["notes"])
        for product in products for tag in product["content_tags"]
    )
    return {
        "requested_count": count,
        "success_count": count,
        "error_count": 0,
        "errors": [],
        "products": products,
        "recovered_from": "xhs_wen_tag_notes",
    }


def cancel_collection(crawler_task_id: str) -> dict[str, Any]:
    """通知爬虫任务在下一个安全检查点停止并保存部分结果。"""

    return _post(f"/api/v1/xhs/wen/tasks/{crawler_task_id}/cancel", {})


def _get(path: str) -> dict[str, Any]:
    url = f"{Config.CRAWLER_BASE_URL}{path}"
    try:
        response = CRAWLER_SESSION.get(url, timeout=Config.CRAWLER_TIMEOUT_SECONDS)
        response.raise_for_status()
        payload = response.json()
    except requests.RequestException as exc:
        raise CrawlerError(f"爬虫服务请求失败：{exc}") from exc
    except ValueError as exc:
        raise CrawlerError("爬虫服务返回的不是有效 JSON") from exc
    if payload.get("code") != 200 or not payload.get("success", False):
        raise CrawlerError(payload.get("msg") or "爬虫服务返回失败")
    data = payload.get("data")
    if not isinstance(data, dict):
        raise CrawlerError("爬虫服务响应缺少 data")
    return data


def get_wen_resources() -> dict[str, Any]:
    """查询当前可用 Cookie 和可立即下发截图的设备。"""
    return _get("/api/v1/xhs/wen/resources")


def list_xhs_wen_xapi_devices() -> list[dict[str, Any]]:
    """查询 Framework `xhs_device_accounts` 中启用的设备清单。

    通过 Framework 的调度配置接口读取，返回每台设备的账号、设备编码与
    可用时间窗口等调度字段，供前端「截图设备列表」直接展示真实设备。
    """
    data = _get("/api/v1/xhs/wen/xapi/devices/settings")
    items = data.get("items")
    return items if isinstance(items, list) else []


def fetch_content_detail(note_url: str) -> dict[str, Any]:
    """按小红书笔记链接读取完整笔记详情。"""
    return _post(
        "/api/v1/fetch_content_detail",
        {"plant": "xhs", "link": note_url},
        retry_transient=True,
    )


def schedule_xapi_screenshots(
    search_queries: list[str], device_count: int | None = None, *,
    trigger_source: str = "acceptance",
    xhs_account_ids: list[str] | None = None,
) -> dict[str, Any]:
    """向手机集群立即下发关键词搜索与截图任务。"""
    queries = normalize_collection_keywords(search_queries)
    account_ids = list(dict.fromkeys(
        str(item or "").strip() for item in (xhs_account_ids or []) if str(item or "").strip()
    ))
    payload: dict[str, Any] = {
        "keywords": queries,
        "execute_now": True,
        "trigger_source": trigger_source,
    }
    if account_ids:
        payload["xhs_account_ids"] = account_ids
    else:
        payload["device_count"] = int(device_count or 0)
    return _post(
        "/api/v1/xhs/wen/xapi/schedule",
        payload,
        retry_transient=False,
    )


def create_xapi_device(data: dict[str, Any]) -> dict[str, Any]:
    """在 framework-db 中新增或重新启用人工维护的 XAPI 设备。"""
    import pymysql
    account_id = str(data.get("xhs_account_id") or "").strip()
    if not account_id:
        raise CrawlerError("xhs_account_id 不能为空")
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    record_id = "manual-" + hashlib.sha256(account_id.encode()).hexdigest()[:24]
    db = pymysql.connect(
        host=Config.CRAWLER_DATABASE_HOST, port=Config.CRAWLER_DATABASE_PORT,
        user=Config.CRAWLER_DATABASE_USER, password=Config.CRAWLER_DATABASE_PASSWORD,
        database=Config.FRAMEWORK_DATABASE_NAME, charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor, autocommit=True,
    )
    try:
        with db.cursor() as cursor:
            cursor.execute(
                """INSERT INTO xhs_device_accounts
                   (feishu_record_id,base_token,table_id,view_id,xhs_account_id,
                    account_name,phone,device_code,text_value,active,raw_json,
                    first_created_at,last_synced_at,updated_at)
                   VALUES (%s,'manual','manual',NULL,%s,%s,%s,%s,'人工添加',1,%s,%s,%s,%s)
                   ON DUPLICATE KEY UPDATE account_name=VALUES(account_name),
                    phone=VALUES(phone),device_code=VALUES(device_code),active=1,
                    raw_json=VALUES(raw_json),last_synced_at=VALUES(last_synced_at),
                    updated_at=VALUES(updated_at)""",
                (record_id, account_id, str(data.get("account_name") or "").strip() or None,
                 str(data.get("phone") or "").strip() or None,
                 str(data.get("device_code") or "").strip() or None,
                 json.dumps(data, ensure_ascii=False), now, now, now),
            )
    finally:
        db.close()
    return {"xhs_account_id": account_id, "active": True}


def delete_xapi_device(account_id: str) -> dict[str, Any]:
    """软删除设备；历史任务保留，后续下发不再选择该设备。"""
    import pymysql
    normalized = str(account_id or "").strip()
    if not normalized:
        raise CrawlerError("xhs_account_id 不能为空")
    db = pymysql.connect(
        host=Config.CRAWLER_DATABASE_HOST, port=Config.CRAWLER_DATABASE_PORT,
        user=Config.CRAWLER_DATABASE_USER, password=Config.CRAWLER_DATABASE_PASSWORD,
        database=Config.FRAMEWORK_DATABASE_NAME, charset="utf8mb4", autocommit=True,
    )
    try:
        with db.cursor() as cursor:
            cursor.execute(
                "UPDATE xhs_device_accounts SET active=0,updated_at=NOW() WHERE xhs_account_id=%s AND active=1",
                (normalized,),
            )
            if not cursor.rowcount:
                raise CrawlerError("设备不存在或已经删除")
    finally:
        db.close()
    return {"xhs_account_id": normalized, "active": False}


def sync_dashboard_acceptance_project(
    batch_key: str, name: str, keywords: list[str], match_keywords: list[str],
) -> None:
    """把验收项目配置同步到 V2；问题和匹配词均按主体隔离。"""
    import pymysql
    normalized_key = str(batch_key or "").strip()
    subject_name = str(name or "").strip() or "未命名监测项目"
    if not normalized_key:
        return
    db = pymysql.connect(
        host=Config.CRAWLER_DATABASE_HOST, port=Config.CRAWLER_DATABASE_PORT,
        user=Config.CRAWLER_DATABASE_USER, password=Config.CRAWLER_DATABASE_PASSWORD,
        database=Config.FRAMEWORK_DATABASE_NAME, charset="utf8mb4",
        cursorclass=pymysql.cursors.DictCursor, autocommit=False,
    )
    try:
        with db.cursor() as cursor:
            cursor.execute("INSERT INTO monitor_subjects(name) VALUES(%s) ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id),status='active'",(subject_name,))
            subject_id = cursor.lastrowid
            cursor.execute("SELECT id FROM monitor_subjects WHERE id=%s FOR UPDATE",(subject_id,))
            cursor.execute("SELECT id FROM monitor_batches WHERE legacy_batch_key=%s AND subject_id=%s AND status<>'archived' ORDER BY collected_at DESC LIMIT 1 FOR UPDATE",(normalized_key,subject_id))
            batch = cursor.fetchone()
            if batch:
                batch_id=batch['id']
            else:
                collection_key='acceptance-'+hashlib.sha256((normalized_key+subject_name).encode()).hexdigest()[:24]
                cursor.execute("INSERT INTO monitor_batches(subject_id,batch_key,legacy_batch_key,name,collected_at,status) VALUES(%s,%s,%s,%s,NOW(),'completed')",(subject_id,collection_key,normalized_key,subject_name))
                batch_id=cursor.lastrowid
            cursor.execute("UPDATE monitor_batch_questions SET status='archived' WHERE batch_id=%s",(batch_id,))
            scopes=[]
            for index, keyword in enumerate(dict.fromkeys(str(v).strip() for v in keywords if str(v).strip()),1):
                cursor.execute("INSERT INTO monitor_questions(subject_id,question_text,display_order) VALUES(%s,%s,%s) ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id),display_order=VALUES(display_order),status='active'",(subject_id,keyword,index))
                question_id=cursor.lastrowid
                cursor.execute("""INSERT INTO monitor_batch_questions(subject_id,batch_id,question_id,question_text_snapshot,display_order,status)
                    VALUES(%s,%s,%s,%s,%s,'active') ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id),question_text_snapshot=VALUES(question_text_snapshot),display_order=VALUES(display_order),status='active'""",(subject_id,batch_id,question_id,keyword,index))
                scopes.append(cursor.lastrowid)
            normalized_matches=[str(v).strip() for v in match_keywords if str(v).strip()]
            for scope_id in scopes:
                cursor.execute("UPDATE monitor_batch_question_keywords SET status='archived' WHERE batch_question_id=%s",(scope_id,))
                for order,keyword in enumerate(dict.fromkeys(normalized_matches),1):
                    cursor.execute("INSERT INTO monitor_keywords(subject_id,keyword_text) VALUES(%s,%s) ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID(id),status='active'",(subject_id,keyword))
                    keyword_id=cursor.lastrowid
                    cursor.execute("""INSERT INTO monitor_batch_question_keywords(subject_id,batch_question_id,keyword_id,keyword_text_snapshot,match_method,display_order,status)
                        VALUES(%s,%s,%s,%s,'substring',%s,'active') ON DUPLICATE KEY UPDATE keyword_text_snapshot=VALUES(keyword_text_snapshot),match_method='substring',display_order=VALUES(display_order),status='active'""",(subject_id,scope_id,keyword_id,keyword,order))
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


def query_xapi_batch_runs(batch_id: str, *, refresh: bool = True) -> dict[str, Any]:
    """读取某个 Framework 调度批次下的设备任务、截图与飞书结果。"""
    normalized = str(batch_id or "").strip()
    if not normalized:
        raise CrawlerError("batch_id 不能为空")
    return _get(
        "/api/v1/xhs/wen/xapi/runs"
        f"?batch_id={normalized}&refresh={'true' if refresh else 'false'}"
    )


def collect_keyword_sources(
    keyword: str, credential_id: str | None = None, cookie: str | None = None,
) -> dict[str, Any]:
    """使用指定 Cookie 采集单个关键词。"""
    normalized = normalize_collection_keywords([keyword])[0]
    body: dict[str, Any] = {"keyword": normalized}
    if cookie:
        body["cookie"] = cookie
    elif credential_id:
        body["credential_id"] = credential_id
    return _post(
        "/api/v1/xhs/wen/sources",
        body,
        retry_transient=False,
    )


def match_feishu_notes(keyword: str, task_limit: int = 100) -> dict[str, Any]:
    """调用爬虫服务，比对飞书笔记和问一问采集来源。"""
    return _post(
        "/api/v1/xhs/feishu-notes/matches",
        {"keyword": keyword, "task_limit": task_limit},
        retry_transient=True,
    )


def resume_collection(crawler_task_id: str) -> dict[str, Any]:
    """沿用原爬虫任务 ID，从 framework 保存的检查点继续采集。"""

    return _post(f"/api/v1/xhs/wen/tasks/{crawler_task_id}/resume", {})