From ca8611113eb119e85f3f95752e14b76d48b11102 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E6=B8=85?= Date: Wed, 19 Aug 2026 16:23:39 +0800 Subject: [PATCH] fix: preserve production task and auth integrations --- app.py | 92 +++++++++++++---- config.py | 9 ++ crawler_client.py | 257 ++++++++++++++++++++++++++++++++++++++-------- database.py | 32 ++++++ requirements.txt | 1 + 5 files changed, 326 insertions(+), 65 deletions(-) diff --git a/app.py b/app.py index 222442d..5587ab5 100644 --- a/app.py +++ b/app.py @@ -6,7 +6,9 @@ from flask import Flask, g, jsonify, redirect, request, send_file from auth_service import ( build_dingtalk_login_url, + create_agent_access_token, exchange_dingtalk_user, + get_user_by_agent_token, get_user_by_session_token, revoke_session, upsert_user_and_session, @@ -32,10 +34,12 @@ from task_service import ( TaskError, cancel_task, create_task, + delete_task, get_file, get_task, get_task_extracted, list_tasks, + merge_tasks, recollect_task, resume_task, resume_unfinished_tasks, @@ -50,12 +54,8 @@ def create_app() -> Flask: if not Config.LOCAL_AUTH_BYPASS or request.remote_addr not in {"127.0.0.1", "::1"}: return None return { - "id": "local-dev", - "name": "本地开发用户", - "account": "local-dev", - "avatar": "", - "union_id": "", - "open_id": "", + "id": "local-dev", "name": "本地开发用户", "account": "local-dev", + "avatar": "", "union_id": "", "open_id": "", } @app.before_request @@ -67,9 +67,16 @@ def create_app() -> Flask: if request.path == "/health" or any(request.path.startswith(path) for path in open_paths): return None user = local_dev_user() or get_user_by_session_token(request.cookies.get(Config.SESSION_COOKIE_NAME)) + if user is None: + authorization = request.headers.get("Authorization", "") + if authorization.lower().startswith("bearer "): + user = get_user_by_agent_token(authorization[7:].strip()) if user is None and request.path.startswith("/api/"): return jsonify({ - "error": {"code": "UNAUTHORIZED", "message": "请先通过钉钉授权登录"}, + "error": { + "code": "UNAUTHORIZED", + "message": "登录或 AI 访问 Token 已失效,请重新登录后复制 Link AI", + }, "request_id": g.request_id, }), 401 g.current_user = user @@ -104,10 +111,14 @@ def create_app() -> Flask: @app.route("/api/v1/auth/dingtalk/login", methods=["GET"]) def dingtalk_login(): - return jsonify({ + response = jsonify({ "data": build_dingtalk_login_url(request.args.get("redirect", "/tasks")), "request_id": g.request_id, }) + # OAuth URL 带一次性 state,且 scope/corpId 配置可能更新,禁止缓存。 + response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate" + response.headers["Pragma"] = "no-cache" + return response @app.route("/api/v1/auth/dingtalk/callback", methods=["GET"]) def dingtalk_callback(): @@ -118,8 +129,13 @@ def create_app() -> Flask: frontend_redirect = state.split(":", 1)[1] or frontend_redirect if not code: return redirect(f"{Config.FRONTEND_BASE_URL}/login?error={quote('钉钉授权缺少 code')}") - user_payload = exchange_dingtalk_user(code) - user, token, expires_at = upsert_user_and_session(user_payload) + try: + user_payload = exchange_dingtalk_user(code) + user, token, expires_at = upsert_user_and_session(user_payload) + except TaskError as exc: + return redirect( + f"{Config.FRONTEND_BASE_URL}/login?error={quote(exc.message)}" + ) response = redirect(f"{Config.FRONTEND_BASE_URL}{frontend_redirect}") response.set_cookie( Config.SESSION_COOKIE_NAME, @@ -130,17 +146,6 @@ def create_app() -> Flask: samesite=Config.SESSION_COOKIE_SAMESITE, path="/", ) - # 给前端“Link AI”复制提示词使用。后端鉴权仍读取 hc_session; - # 前端只读取这个辅助 Cookie 的 value,并在提示词中拼成 hc_session=。 - response.set_cookie( - "hc_session_agent", - token, - expires=expires_at, - httponly=False, - secure=Config.SESSION_COOKIE_SECURE, - samesite=Config.SESSION_COOKIE_SAMESITE, - path="/", - ) app.logger.info("钉钉用户登录成功 user_id=%s name=%s", user["id"], user["name"]) return response @@ -216,11 +221,34 @@ def create_app() -> Flask: raise TaskError("UNAUTHORIZED", "请先通过钉钉授权登录", 401) return jsonify({"data": {"user": user}, "request_id": g.request_id}) + @app.route("/api/v1/auth/agent-token", methods=["POST"]) + def auth_agent_token(): + issued = create_agent_access_token( + request.cookies.get(Config.SESSION_COOKIE_NAME) + ) + if not issued: + raise TaskError( + "UNAUTHORIZED", + "登录已失效,请重新完成钉钉登录后再复制 Link AI", + 401, + ) + token, expires_at = issued + response = jsonify({ + "data": { + "token": token, + "expires_at": expires_at.isoformat(timespec="seconds"), + }, + "request_id": g.request_id, + }) + response.headers["Cache-Control"] = "no-store" + return response + @app.route("/api/v1/auth/logout", methods=["POST"]) def auth_logout(): revoke_session(request.cookies.get(Config.SESSION_COOKIE_NAME)) response = jsonify({"data": {"ok": True}, "request_id": g.request_id}) response.delete_cookie(Config.SESSION_COOKIE_NAME, path="/") + # 清理旧版本曾设置的可读辅助 Cookie。 response.delete_cookie("hc_session_agent", path="/") return response @@ -272,12 +300,27 @@ def create_app() -> Flask: ) return jsonify({"data": result, "request_id": g.request_id}) + @app.route("/api/v1/tasks/merge", methods=["POST"]) + def merge_research_tasks(): + payload = request.get_json(silent=True) or {} + task_ids = payload.get("task_ids") + if not isinstance(task_ids, list): + raise TaskError("MERGE_TASKS_REQUIRED", "task_ids 必须是任务 ID 数组", 422) + return jsonify({ + "data": merge_tasks(task_ids), + "request_id": g.request_id, + }) + @app.route("/api/v1/tasks/", methods=["GET"]) def get_research_task(task_id: str): if request.args.get("view") in {"extracted", "cleaned", "parsed"}: return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id}) return jsonify({"data": get_task(task_id), "request_id": g.request_id}) + @app.route("/api/v1/tasks/", methods=["DELETE"]) + def delete_research_task(task_id: str): + return jsonify({"data": delete_task(task_id), "request_id": g.request_id}) + @app.route("/api/v1/tasks//extracted", methods=["GET"]) def get_research_task_extracted(task_id: str): return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id}) @@ -297,8 +340,13 @@ def create_app() -> Flask: @app.route("/api/v1/tasks//recollect", methods=["POST"]) def recollect_research_task(task_id: str): + current_user = getattr(g, "current_user", None) or {} return jsonify({ - "data": recollect_task(task_id), + "data": recollect_task( + task_id, + created_by_id=current_user.get("id"), + created_by_name=current_user.get("name"), + ), "request_id": g.request_id, }), 201 diff --git a/config.py b/config.py index 26dda81..7d3fae8 100644 --- a/config.py +++ b/config.py @@ -35,10 +35,19 @@ class Config: LOCAL_AUTH_BYPASS = os.getenv("LOCAL_AUTH_BYPASS", "false").lower() in ("1", "true", "yes", "on") DINGTALK_CLIENT_ID = os.getenv("DINGTALK_CLIENT_ID", os.getenv("DINGTALK_APP_KEY", "")) DINGTALK_CLIENT_SECRET = os.getenv("DINGTALK_CLIENT_SECRET", os.getenv("DINGTALK_APP_SECRET", "")) + DINGTALK_ALLOWED_CORP_ID = os.getenv( + "DINGTALK_ALLOWED_CORP_ID", + "dingbc8e63f8aa03221d35c2f4657eb6378f", + ) DINGTALK_REDIRECT_URI = os.getenv( "DINGTALK_REDIRECT_URI", f"{FRONTEND_BASE_URL}/api/v1/auth/dingtalk/callback", ) + CRAWLER_DATABASE_HOST = os.getenv("CRAWLER_DATABASE_HOST", "") + CRAWLER_DATABASE_PORT = int(os.getenv("CRAWLER_DATABASE_PORT", "3306")) + CRAWLER_DATABASE_USER = os.getenv("CRAWLER_DATABASE_USER", "") + CRAWLER_DATABASE_PASSWORD = os.getenv("CRAWLER_DATABASE_PASSWORD", "") + CRAWLER_DATABASE_NAME = os.getenv("CRAWLER_DATABASE_NAME", "") DINGTALK_AUTH_URL = os.getenv("DINGTALK_AUTH_URL", "https://login.dingtalk.com/oauth2/auth") DINGTALK_TOKEN_URL = os.getenv("DINGTALK_TOKEN_URL", "https://api.dingtalk.com/v1.0/oauth2/userAccessToken") DINGTALK_USER_URL = os.getenv("DINGTALK_USER_URL", "https://api.dingtalk.com/v1.0/contact/users/me") diff --git a/crawler_client.py b/crawler_client.py index d163757..33a092e 100644 --- a/crawler_client.py +++ b/crawler_client.py @@ -1,8 +1,10 @@ import re +import secrets import time from typing import Any import requests +import json from config import Config @@ -11,6 +13,10 @@ class CrawlerError(RuntimeError): pass +class NoAiCardError(CrawlerError): + """搜索请求成功,但当前结果中没有问一问 AI 卡片。""" + + CRAWLER_SESSION = requests.Session() # 爬虫服务为明确配置的直连地址,不应继承 macOS 系统代理。 # 否则 requests 会经由 127.0.0.1:7890,约 30 秒后被代理返回 502。 @@ -30,6 +36,58 @@ def normalize_collection_keywords(keywords: list[str]) -> list[str]: 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], @@ -71,13 +129,160 @@ def _post( raise CrawlerError(f"爬虫服务请求失败(已重试 {retry_times} 次)") if payload.get("code") != 200 or not payload.get("success", False): - raise CrawlerError(payload.get("msg") or "爬虫服务返回失败") + 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: @@ -101,46 +306,8 @@ def get_wen_resources() -> dict[str, Any]: return _get("/api/v1/xhs/wen/resources") -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 只有接收数组时才创建异步任务并立即返回 task_id。 - # 因此轻度任务即使只有一个关键词也必须传数组,避免单关键词走同步 - # overview 并因采集耗时超过网关时限而返回 502。 - keyword: str | list[str] = ( - collection_keywords if mode == "mild" else collection_keywords[0] - ) - # 轻度 overview 在首次计算超过网关时限后,远程通常已经写好缓存, - # 允许受控重试以取得结果;深度 sources 会启动长任务,禁止盲目重试, - # 避免远程请求虽断开但后台继续执行时产生重复采集任务。 - return _post( - path, - {"keyword": keyword}, - retry_transient=(mode != "deep"), - ) - - -def query_collection(crawler_task_id: str) -> dict[str, Any]: - return _post("/api/v1/xhs/wen/tasks/query", {"task_id": crawler_task_id}) - - -def cancel_collection(crawler_task_id: str) -> dict[str, Any]: - """通知爬虫任务在下一个安全检查点停止并保存部分结果。""" - - return _post(f"/api/v1/xhs/wen/tasks/{crawler_task_id}/cancel", {}) - - -def resume_collection(crawler_task_id: str) -> dict[str, Any]: - """沿用原爬虫任务 ID,从 framework 保存的检查点继续采集。""" - - return _post(f"/api/v1/xhs/wen/tasks/{crawler_task_id}/resume", {}) - - def fetch_content_detail(note_url: str) -> dict[str, Any]: """按小红书笔记链接读取完整笔记详情。""" - return _post( "/api/v1/fetch_content_detail", {"plant": "xhs", "link": note_url}, @@ -149,8 +316,7 @@ def fetch_content_detail(note_url: str) -> dict[str, Any]: def schedule_xapi_screenshots(search_queries: list[str], device_count: int) -> dict[str, Any]: - """向手机集群下发问一问关键词搜索与截图任务。""" - + """向手机集群下发关键词搜索与截图任务。""" queries = normalize_collection_keywords(search_queries) return _post( "/api/v1/xhs/wen/xapi/schedule", @@ -160,8 +326,7 @@ def schedule_xapi_screenshots(search_queries: list[str], device_count: int) -> d def collect_keyword_sources(keyword: str, credential_id: str | None = None) -> dict[str, Any]: - """使用 sources 的完整 Cookie 池采集单个关键词。""" - + """使用指定 Cookie 采集单个关键词。""" normalized = normalize_collection_keywords([keyword])[0] body: dict[str, Any] = {"keyword": normalized} if credential_id: @@ -171,3 +336,9 @@ def collect_keyword_sources(keyword: str, credential_id: str | None = None) -> d body, retry_transient=False, ) + + +def resume_collection(crawler_task_id: str) -> dict[str, Any]: + """沿用原爬虫任务 ID,从 framework 保存的检查点继续采集。""" + + return _post(f"/api/v1/xhs/wen/tasks/{crawler_task_id}/resume", {}) diff --git a/database.py b/database.py index 7f23510..34dc98f 100644 --- a/database.py +++ b/database.py @@ -85,6 +85,25 @@ def init_database() -> None: FOREIGN KEY(task_id) REFERENCES research_tasks(id) ON DELETE CASCADE ); + CREATE TABLE IF NOT EXISTS task_keyword_runs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + task_id INTEGER NOT NULL, + keyword_order INTEGER NOT NULL, + keyword TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'queued', + progress_percent INTEGER NOT NULL DEFAULT 0, + crawler_task_id TEXT, + result_json TEXT, + error_message TEXT, + started_at TEXT, + completed_at TEXT, + updated_at TEXT NOT NULL, + UNIQUE(task_id, keyword_order), + FOREIGN KEY(task_id) REFERENCES research_tasks(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_keyword_runs_task_status + ON task_keyword_runs(task_id, status, keyword_order); + CREATE TABLE IF NOT EXISTS task_files ( id INTEGER PRIMARY KEY AUTOINCREMENT, public_file_id TEXT NOT NULL UNIQUE, @@ -134,6 +153,19 @@ def init_database() -> None: CREATE INDEX IF NOT EXISTS idx_auth_sessions_user ON auth_sessions(user_id, expires_at); + CREATE TABLE IF NOT EXISTS ai_access_tokens ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + token_hash TEXT NOT NULL UNIQUE, + user_id INTEGER NOT NULL, + expires_at TEXT NOT NULL, + created_at TEXT NOT NULL, + last_seen_at TEXT NOT NULL, + revoked_at TEXT, + FOREIGN KEY(user_id) REFERENCES auth_users(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_ai_access_tokens_user + ON ai_access_tokens(user_id, expires_at); + CREATE TABLE IF NOT EXISTS acceptance_tasks ( id INTEGER PRIMARY KEY AUTOINCREMENT, public_task_id TEXT NOT NULL UNIQUE, diff --git a/requirements.txt b/requirements.txt index 4fb9e89..e7075c7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ Flask==3.1.1 requests==2.32.4 python-dotenv==1.1.1 gunicorn==23.0.0 +PyMySQL==1.1.1 playwright==1.58.0