"""430 query 真实采集:搜索前 10、详情限速、媒体保存/OSS、创作知识判断。""" from __future__ import annotations import hashlib import json import time from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Iterable, Optional from acquisition import store from acquisition.classify import classify_imgtext, classify_video from acquisition.crawler import RateLimiter, fetch_post_detail from acquisition.oss import upload_stream from acquisition.search import search_keyword, search_weixin, search_xiaohongshu, fetch_weixin_detail from core.config import Settings from core.models import Post from creation_knowledge.integrations.video_extract import _default_download ROOT = Path(__file__).resolve().parent.parent DEFAULT_QUERY_FILE = ROOT / "data" / "queries" / "creation_demo.json" PLATFORMS = ("xiaohongshu", "weixin", "douyin") Downloader = Callable[[str, str], bytes] class PlatformRateLimiter: """把同一平台的 keyword/detail 调用合并到同一个 10-12s bucket。""" def __init__(self, platform: str, *, min_seconds: float = 10.0, max_seconds: float = 12.0, delegate: Optional[RateLimiter] = None) -> None: self.platform = platform self.delegate = delegate or RateLimiter( min_interval_seconds=min_seconds, max_interval_seconds=max_seconds, ) def wait(self, bucket: str) -> None: self.delegate.wait(self.platform) @dataclass class Candidate: rank: int source_id: str = "" url: str = "" title: str = "" author: str = "" cover_url: str = "" raw: dict | None = None def load_creation_queries(path: Path | str = DEFAULT_QUERY_FILE) -> list[str]: """从 creation_demo.json 读取所有 keep=true 的唯一 query,保持首次出现顺序。""" data = json.loads(Path(path).read_text("utf-8")) seen: set[str] = set() out: list[str] = [] for fam in data.get("families") or []: for item in fam.get("items") or []: q = (item.get("query") or "").strip() if q and item.get("keep", True) and q not in seen: seen.add(q) out.append(q) return out def prompt_version(*names: str) -> str: h = hashlib.sha256() for name in names: p = ROOT / "prompts" / f"{name}.txt" h.update(name.encode("utf-8")) if p.exists(): h.update(p.read_bytes()) return h.hexdigest()[:16] def _hash(value: str, n: int = 16) -> str: return hashlib.md5(value.encode("utf-8")).hexdigest()[:n] def _ext_from_bytes(data: bytes) -> str: if data[:4] == b"\x89PNG": return ".png" if data[:3] == b"GIF": return ".gif" if data[:4] == b"RIFF" and data[8:12] == b"WEBP": return ".webp" return ".jpg" def _dedup(values: Iterable[str]) -> list[str]: seen: set[str] = set() out: list[str] = [] for v in values: if v and v not in seen: seen.add(v) out.append(v) return out def _data_root(settings: Settings) -> Path: p = Path(settings.data_dir or "data") return p if p.is_absolute() else ROOT / p def _save_images(urls: list[str], *, platform: str, base_dir: Path, public_base: str, downloader: Downloader = _default_download, max_images: int = 12) -> list[str]: out: list[str] = [] for i, url in enumerate(_dedup(urls)[:max_images], start=1): if not url.startswith("http"): continue try: data = downloader(url, platform) except Exception: continue if not data: continue base_dir.mkdir(parents=True, exist_ok=True) ext = _ext_from_bytes(data) name = f"image_{i}{ext}" (base_dir / name).write_bytes(data) out.append(f"{public_base}/{name}") return out def _classify_and_store(conn, item_id: int, platform: str, *, title: str, body_text: str, images: list[str], video_url: str, settings: Settings, ts: Optional[int] = None) -> None: if platform == "douyin": version = prompt_version("classify_video") is_creation, reason, knowledge, _points = classify_video( {"platform": platform, "title": title, "body_text": body_text, "video": video_url}, settings, ) else: version = prompt_version("classify_imgtext") is_creation, reason, knowledge, _points = classify_imgtext( {"platform": platform, "title": title, "body_text": body_text, "images": images}, settings, ) store.upsert_creation_classification( conn, item_id, is_creation, reason=reason, knowledge=knowledge, prompt_version=version, error=reason if is_creation is None else "", ts=ts, ) def search_candidates(platform: str, query: str, *, settings: Settings, limit: int, rate_limiter: PlatformRateLimiter) -> list[Candidate]: if platform == "xiaohongshu": rows = search_xiaohongshu( query, content_type="图文", limit=limit, settings=settings, rate_limiter=rate_limiter ) return [ Candidate( rank=i, source_id=r.get("id") or "", url=r.get("url") or "", title=r.get("title") or "", author=r.get("nick_name") or "", cover_url=r.get("cover_url") or "", raw=r, ) for i, r in enumerate(rows, start=1) ] if platform == "weixin": rows = search_weixin(query, limit=limit, settings=settings, rate_limiter=rate_limiter) return [ Candidate( rank=i, url=r.get("url") or "", title=r.get("title") or "", author=r.get("nick_name") or "", cover_url=r.get("cover_url") or "", raw=r, ) for i, r in enumerate(rows, start=1) ] if platform == "douyin": ids = search_keyword( query, platform="douyin", content_type="视频", limit=limit, settings=settings, rate_limiter=rate_limiter, ) return [Candidate(rank=i, source_id=cid, raw={"id": cid}) for i, cid in enumerate(ids, start=1)] raise ValueError(f"unsupported platform: {platform}") def _process_xhs(candidate: Candidate, *, query_hash: str, settings: Settings, rate_limiter: PlatformRateLimiter, downloader: Downloader) -> dict: post = fetch_post_detail(candidate.source_id or candidate.url, settings=settings, rate_limiter=rate_limiter) base = _data_root(settings) / "media" / "xiaohongshu" / query_hash / (post.content_id or candidate.source_id) pub = f"/data/media/xiaohongshu/{query_hash}/{post.content_id or candidate.source_id}" images = _save_images(post.image_urls or [candidate.cover_url], platform="xiaohongshu", base_dir=base, public_base=pub, downloader=downloader) if not images: raise RuntimeError("图片下载失败") return { "source_id": post.content_id or candidate.source_id, "url": post.url or candidate.url, "title": post.title or candidate.title, "author": post.author_name or candidate.author, "body_text": post.body_text or "", "cover_url": images[0], "image_urls": images, "video_url": "", "raw": post.raw if isinstance(post.raw, dict) else {}, } def _process_weixin(candidate: Candidate, *, query_hash: str, settings: Settings, rate_limiter: PlatformRateLimiter, downloader: Downloader) -> dict: body_text, detail_images = fetch_weixin_detail( candidate.url, settings=settings, rate_limiter=rate_limiter ) h = _hash(candidate.url) base = _data_root(settings) / "media" / "weixin" / query_hash / h pub = f"/data/media/weixin/{query_hash}/{h}" images = _save_images(detail_images or [candidate.cover_url], platform="weixin", base_dir=base, public_base=pub, downloader=downloader) return { "source_id": h, "url": candidate.url, "title": candidate.title, "author": candidate.author, "body_text": body_text, "cover_url": images[0] if images else candidate.cover_url, "image_urls": images, "video_url": "", "raw": candidate.raw or {}, } def _process_douyin(candidate: Candidate, *, settings: Settings, rate_limiter: PlatformRateLimiter) -> dict: post: Post = fetch_post_detail(candidate.source_id, settings=settings, rate_limiter=rate_limiter) if not post.video_urls: raise RuntimeError("详情无视频") cdn_url = upload_stream(post.video_urls[0], src_type="video", settings=settings) return { "source_id": post.content_id or candidate.source_id, "url": post.url, "title": post.title, "author": post.author_name or "", "body_text": post.body_text or "", "cover_url": (post.image_urls or [""])[0], "image_urls": post.image_urls or [], "video_url": cdn_url, "raw": post.raw if isinstance(post.raw, dict) else {}, } def process_candidate(platform: str, candidate: Candidate, *, query_hash: str, settings: Settings, rate_limiter: PlatformRateLimiter, downloader: Downloader = _default_download) -> dict: if platform == "xiaohongshu": return _process_xhs(candidate, query_hash=query_hash, settings=settings, rate_limiter=rate_limiter, downloader=downloader) if platform == "weixin": return _process_weixin(candidate, query_hash=query_hash, settings=settings, rate_limiter=rate_limiter, downloader=downloader) if platform == "douyin": return _process_douyin(candidate, settings=settings, rate_limiter=rate_limiter) raise ValueError(f"unsupported platform: {platform}") def run_platform_query(conn, *, run_id: str, query: str, platform: str, settings: Settings, search_limit: int = 10, display_limit: int = 5, rate_limiter: Optional[PlatformRateLimiter] = None, downloader: Downloader = _default_download, sleep_fn: Callable[[float], None] = time.sleep, classify: bool = True, ts: Optional[int] = None) -> dict: """执行一个 query×platform job。任何 item 失败都会落库并继续。""" gate = rate_limiter or PlatformRateLimiter(platform) now = int(time.time()) if ts is None else ts store.update_creation_job(conn, run_id, query, platform, status="running", ts=now) candidates: list[Candidate] = [] last_error = "" attempts = 0 for attempts in range(1, 4): try: candidates = search_candidates( platform, query, settings=settings, limit=search_limit, rate_limiter=gate ) last_error = "" break except Exception as exc: last_error = f"搜索失败: {str(exc)[:120]}" store.update_creation_job( conn, run_id, query, platform, status="running", attempts=attempts, error=last_error, ts=int(time.time()), ) if attempts < 3: sleep_fn(30 if attempts == 1 else 90) if last_error: store.update_creation_job( conn, run_id, query, platform, status="failed", attempts=attempts, searched_count=0, display_count=0, error=last_error, ts=int(time.time()), ) return {"platform": platform, "status": "failed", "display_count": 0, "error": last_error} qh = _hash(query) display_count = 0 errors: list[str] = [] for cand in candidates[:search_limit]: if display_count >= display_limit: break try: processed = process_candidate( platform, cand, query_hash=qh, settings=settings, rate_limiter=gate, downloader=downloader, ) item_id = store.upsert_creation_item( conn, run_id=run_id, query=query, platform=platform, rank=cand.rank, source_id=processed.get("source_id", ""), url=processed.get("url", ""), title=processed.get("title", ""), author=processed.get("author", ""), body_text=processed.get("body_text", ""), cover_url=processed.get("cover_url", ""), image_urls=processed.get("image_urls", []), video_url=processed.get("video_url", ""), media={"source": cand.raw or {}, "candidate_rank": cand.rank}, raw=processed.get("raw", {}), is_displayable=True, ts=int(time.time()), ) if classify: _classify_and_store( conn, item_id, platform, title=processed.get("title", ""), body_text=processed.get("body_text", ""), images=processed.get("image_urls", []), video_url=processed.get("video_url", ""), settings=settings, ts=int(time.time()), ) display_count += 1 except Exception as exc: msg = f"rank {cand.rank}: {str(exc)[:120]}" errors.append(msg) store.upsert_creation_item( conn, run_id=run_id, query=query, platform=platform, rank=cand.rank, source_id=cand.source_id, url=cand.url, title=cand.title, author=cand.author, cover_url=cand.cover_url, raw=cand.raw or {}, is_displayable=False, error=msg, ts=int(time.time()), ) if display_count >= display_limit: status = "done" error = "" elif display_count > 0: status = "partial" error = "; ".join(errors[-3:]) else: status = "failed" error = "; ".join(errors[-3:]) or "未获得可展示内容" store.update_creation_job( conn, run_id, query, platform, status=status, attempts=attempts, searched_count=len(candidates), display_count=display_count, error=error, ts=int(time.time()), ) return {"platform": platform, "status": status, "display_count": display_count, "error": error}