Add Cloud Tour Libo knowledge graph platform
This commit is contained in:
294
scripts/douyin_fetch_data.py
Normal file
294
scripts/douyin_fetch_data.py
Normal file
@@ -0,0 +1,294 @@
|
||||
#!/usr/bin/env python3
|
||||
"""单独拉取抖音数据的运行入口。
|
||||
|
||||
这个文件不重写采集逻辑，而是复用系统原版:
|
||||
app/agents/douyin_agent.py
|
||||
|
||||
常用命令:
|
||||
python3 scripts/douyin_login.py
|
||||
python3 scripts/douyin_fetch_data.py --name "某某酒店" --city "贵阳"
|
||||
python3 scripts/douyin_fetch_data.py --input hotels.txt --city "贵阳" --save-evidence
|
||||
|
||||
输出:
|
||||
默认写到 data/douyin_probe/fetch_<时间>.json
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import csv
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
from app.agents.douyin_agent import ( # noqa: E402
|
||||
_collect,
|
||||
_parse_dy_comments,
|
||||
_parse_dy_notes,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_OUT_DIR = ROOT / "data" / "douyin_probe"
|
||||
|
||||
|
||||
def _ts() -> str:
|
||||
return datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
|
||||
|
||||
def _read_names(path: str | None) -> list[str]:
|
||||
if not path:
|
||||
return []
|
||||
names: list[str] = []
|
||||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||||
value = line.strip()
|
||||
if value and not value.startswith("#"):
|
||||
names.append(value)
|
||||
return names
|
||||
|
||||
|
||||
def _dedupe_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
deduped: list[dict[str, Any]] = []
|
||||
seen: set[tuple[str, str, str]] = set()
|
||||
for record in records:
|
||||
key = (
|
||||
str(record.get("platform") or ""),
|
||||
str(record.get("kind") or ""),
|
||||
str(record.get("source_id") or record.get("url") or ""),
|
||||
)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
deduped.append(record)
|
||||
return deduped
|
||||
|
||||
|
||||
def _record_for_csv(record: dict[str, Any]) -> dict[str, Any]:
|
||||
fields = [
|
||||
"platform",
|
||||
"kind",
|
||||
"source_id",
|
||||
"url",
|
||||
"entity_name",
|
||||
"place_natural_key",
|
||||
"keyword",
|
||||
"title",
|
||||
"content",
|
||||
"author",
|
||||
"author_id",
|
||||
"likes",
|
||||
"comments",
|
||||
"collects",
|
||||
"shares",
|
||||
"publish_time",
|
||||
"location",
|
||||
]
|
||||
return {field: record.get(field, "") for field in fields}
|
||||
|
||||
|
||||
def _write_json(path: Path, payload: dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", encoding="utf-8") as file:
|
||||
for record in records:
|
||||
file.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
|
||||
|
||||
def _write_csv(path: Path, records: list[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
rows = [_record_for_csv(record) for record in records]
|
||||
fieldnames = list(rows[0].keys()) if rows else list(_record_for_csv({}).keys())
|
||||
with path.open("w", encoding="utf-8-sig", newline="") as file:
|
||||
writer = csv.DictWriter(file, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def _build_keyword(name: str, city: str, suffix: str, keyword: str | None) -> str:
|
||||
if keyword:
|
||||
return keyword.strip()
|
||||
parts = [city.strip(), name.strip(), suffix.strip()]
|
||||
return " ".join(part for part in parts if part)
|
||||
|
||||
|
||||
def _make_tab_picker(preferred_tab: str):
|
||||
def pick(labels: list[str]) -> str | None:
|
||||
if not labels:
|
||||
return None
|
||||
if preferred_tab and preferred_tab in labels:
|
||||
return preferred_tab
|
||||
if "综合" in labels:
|
||||
return "综合"
|
||||
return labels[0]
|
||||
|
||||
return pick
|
||||
|
||||
|
||||
def fetch_one(
|
||||
*,
|
||||
name: str,
|
||||
city: str,
|
||||
suffix: str,
|
||||
keyword: str | None,
|
||||
tab: str,
|
||||
deep: bool,
|
||||
place_key: str | None,
|
||||
) -> dict[str, Any]:
|
||||
query = _build_keyword(name, city, suffix, keyword)
|
||||
print(f"[douyin] 开始采集: name={name} keyword={query}")
|
||||
result = _collect(query, _make_tab_picker(tab), deep=deep)
|
||||
|
||||
if result.get("logged_in") is False:
|
||||
return {
|
||||
"ok": False,
|
||||
"need_login": True,
|
||||
"name": name,
|
||||
"keyword": query,
|
||||
"summary": "抖音未登录。请先运行 python3 scripts/douyin_login.py 完成一次登录。",
|
||||
"notes": [],
|
||||
"comments": [],
|
||||
"records": [],
|
||||
"raw_api_count": 0,
|
||||
"api_url_count": result.get("api_url_count", 0),
|
||||
"tabs": result.get("tabs", []),
|
||||
}
|
||||
|
||||
if result.get("error"):
|
||||
return {
|
||||
"ok": False,
|
||||
"need_login": False,
|
||||
"name": name,
|
||||
"keyword": query,
|
||||
"summary": f"采集异常: {result.get('error')}",
|
||||
"notes": [],
|
||||
"comments": [],
|
||||
"records": [],
|
||||
"raw_api_count": 0,
|
||||
"api_url_count": result.get("api_url_count", 0),
|
||||
"tabs": result.get("tabs", []),
|
||||
}
|
||||
|
||||
raw_api = result.get("raw_api") or []
|
||||
notes = _parse_dy_notes(raw_api, name, query)
|
||||
comments = _parse_dy_comments(raw_api, name, query)
|
||||
records = _dedupe_records(notes + comments)
|
||||
for record in records:
|
||||
if place_key:
|
||||
record["place_natural_key"] = place_key
|
||||
|
||||
return {
|
||||
"ok": True,
|
||||
"need_login": False,
|
||||
"name": name,
|
||||
"keyword": query,
|
||||
"summary": f"视频 {len(notes)} 条，评论 {len(comments)} 条，去重后 {len(records)} 条。",
|
||||
"notes": notes,
|
||||
"comments": comments,
|
||||
"records": records,
|
||||
"raw_api_count": len(raw_api),
|
||||
"api_url_count": result.get("api_url_count", 0),
|
||||
"tabs": result.get("tabs", []),
|
||||
}
|
||||
|
||||
|
||||
async def _save_evidence(records: list[dict[str, Any]]) -> int:
|
||||
if not records:
|
||||
return 0
|
||||
from app.db import close_pool, init_pool, sa_save_evidence
|
||||
|
||||
await init_pool()
|
||||
try:
|
||||
return await sa_save_evidence(records)
|
||||
finally:
|
||||
await close_pool()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="按酒店/关键词单独获取抖音视频与评论数据")
|
||||
parser.add_argument("--name", help="单个酒店/地点名称")
|
||||
parser.add_argument("--input", help="批量名称文件，每行一个酒店/地点名称")
|
||||
parser.add_argument("--city", default="贵阳", help="城市前缀，默认: 贵阳")
|
||||
parser.add_argument("--suffix", default="", help="附加搜索词，例如: 酒店 入住 评价")
|
||||
parser.add_argument("--keyword", help="完整搜索关键词；传入后会忽略 --city/--name/--suffix 拼接")
|
||||
parser.add_argument("--tab", default="综合", help="优先选择的抖音搜索 tab，默认: 综合")
|
||||
parser.add_argument("--place-key", help="写入记录的 place_natural_key，单名称时可用")
|
||||
parser.add_argument("--no-deep", action="store_true", help="只采搜索页，不进入视频详情深采评论")
|
||||
parser.add_argument("--save-evidence", action="store_true", help="把 records 写入系统 Evidence 表")
|
||||
parser.add_argument("--out", help="JSON 输出路径")
|
||||
parser.add_argument("--jsonl-out", help="JSONL records 输出路径")
|
||||
parser.add_argument("--csv-out", help="CSV records 输出路径")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
names = []
|
||||
if args.name:
|
||||
names.append(args.name.strip())
|
||||
names.extend(_read_names(args.input))
|
||||
names = [name for name in names if name]
|
||||
|
||||
if args.keyword and not names:
|
||||
names = [args.keyword.strip()]
|
||||
if not names:
|
||||
raise SystemExit("请提供 --name、--input 或 --keyword")
|
||||
if args.keyword and len(names) > 1:
|
||||
raise SystemExit("--keyword 是完整单次搜索词，不能和多名称 --input 同时使用")
|
||||
if args.place_key and len(names) > 1:
|
||||
raise SystemExit("--place-key 只适合单个 --name 使用，批量请不要传")
|
||||
|
||||
all_results: list[dict[str, Any]] = []
|
||||
all_records: list[dict[str, Any]] = []
|
||||
for name in names:
|
||||
result = fetch_one(
|
||||
name=name,
|
||||
city=args.city,
|
||||
suffix=args.suffix,
|
||||
keyword=args.keyword,
|
||||
tab=args.tab,
|
||||
deep=not args.no_deep,
|
||||
place_key=args.place_key,
|
||||
)
|
||||
print(f"[douyin] {result['summary']}")
|
||||
all_results.append(result)
|
||||
all_records.extend(result.get("records") or [])
|
||||
|
||||
all_records = _dedupe_records(all_records)
|
||||
saved = 0
|
||||
if args.save_evidence:
|
||||
saved = asyncio.run(_save_evidence(all_records))
|
||||
print(f"[douyin] 已写入 Evidence: {saved}")
|
||||
|
||||
out_dir = DEFAULT_OUT_DIR
|
||||
out_path = Path(args.out) if args.out else out_dir / f"fetch_{_ts()}.json"
|
||||
jsonl_path = Path(args.jsonl_out) if args.jsonl_out else out_path.with_suffix(".jsonl")
|
||||
csv_path = Path(args.csv_out) if args.csv_out else out_path.with_suffix(".csv")
|
||||
|
||||
payload = {
|
||||
"ok": all(result.get("ok") for result in all_results),
|
||||
"need_login": any(result.get("need_login") for result in all_results),
|
||||
"saved_evidence": saved,
|
||||
"total_records": len(all_records),
|
||||
"results": all_results,
|
||||
}
|
||||
_write_json(out_path, payload)
|
||||
_write_jsonl(jsonl_path, all_records)
|
||||
_write_csv(csv_path, all_records)
|
||||
|
||||
print(f"[douyin] JSON: {out_path}")
|
||||
print(f"[douyin] JSONL: {jsonl_path}")
|
||||
print(f"[douyin] CSV: {csv_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.chdir(ROOT)
|
||||
main()
|
||||
Reference in New Issue
Block a user