store.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694
  1. """SQLite 数据层:query 生成结果 + 真实搜索结果,统一进 data/app.db。
  2. 为什么不再用散落 json:量大了「一次性 fetch 整个 json」撑不住,要能分页/筛选/多 run 累积。
  3. 设计:
  4. queries —— 一条生成的 query 一行;各打法字段不同,统一塞进 axes(json) 列
  5. search_results —— 一条「query × 平台」一行(把 douyin/weixin 嵌套拍平),方便按平台/成功筛
  6. 媒体文件仍躺 data/search/,库里只存相对路径(cover/video)。
  7. 纯 stdlib sqlite3,无新依赖。写入幂等:同 run 重导先清该 run。
  8. """
  9. from __future__ import annotations
  10. import json
  11. import os
  12. import sqlite3
  13. from pathlib import Path
  14. from typing import Any, Iterable, Optional
  15. from core.config import load_env_file
  16. ROOT = Path(__file__).resolve().parent.parent
  17. DB_PATH = ROOT / "data" / "app.db"
  18. _SCHEMA = """
  19. CREATE TABLE IF NOT EXISTS queries (
  20. id INTEGER PRIMARY KEY AUTOINCREMENT,
  21. run_id TEXT NOT NULL,
  22. method TEXT NOT NULL, -- 打法名(实质 × 创作阶段 × 需求 等)
  23. query TEXT NOT NULL,
  24. axes TEXT NOT NULL DEFAULT '{}', -- 该打法的正交轴 json(实质/形式/阶段…)
  25. created_at INTEGER NOT NULL
  26. );
  27. CREATE INDEX IF NOT EXISTS ix_queries_run ON queries(run_id);
  28. CREATE INDEX IF NOT EXISTS ix_queries_method ON queries(method);
  29. CREATE TABLE IF NOT EXISTS search_results (
  30. id INTEGER PRIMARY KEY AUTOINCREMENT,
  31. run_id TEXT NOT NULL,
  32. method TEXT NOT NULL,
  33. query TEXT NOT NULL,
  34. platform TEXT NOT NULL, -- douyin / weixin
  35. ok INTEGER NOT NULL, -- 1=搜到且有媒体, 0=失败/空
  36. title TEXT,
  37. url TEXT,
  38. cover TEXT, -- /data/search/... 相对路径
  39. video TEXT, -- 抖音才有
  40. extra TEXT NOT NULL DEFAULT '{}', -- nick/error 等附加 json
  41. created_at INTEGER NOT NULL
  42. );
  43. CREATE INDEX IF NOT EXISTS ix_sr_run ON search_results(run_id);
  44. CREATE INDEX IF NOT EXISTS ix_sr_platform ON search_results(platform);
  45. CREATE INDEX IF NOT EXISTS ix_sr_method ON search_results(method);
  46. CREATE TABLE IF NOT EXISTS runs (
  47. run_id TEXT PRIMARY KEY,
  48. kind TEXT NOT NULL, -- queries / search
  49. note TEXT,
  50. created_at INTEGER NOT NULL
  51. );
  52. -- 帖子级「创作知识 / 非创作知识」分类(按 url,一帖判一次)。acquisition/classify.py 写入。
  53. CREATE TABLE IF NOT EXISTS post_class (
  54. url TEXT PRIMARY KEY,
  55. is_creation INTEGER, -- 1=创作知识 / 0=非
  56. reason TEXT,
  57. knowledge TEXT, -- 提取出的完整创作知识(创作帖才有)
  58. points TEXT, -- 结构化知识点 json(按卡片/分段)
  59. created_at INTEGER NOT NULL
  60. );
  61. -- 430 query 真实采集:run / query×platform job / item / AI 判断
  62. CREATE TABLE IF NOT EXISTS creation_search_runs (
  63. run_id TEXT PRIMARY KEY,
  64. status TEXT NOT NULL,
  65. total_queries INTEGER NOT NULL DEFAULT 0,
  66. note TEXT,
  67. started_at INTEGER NOT NULL,
  68. finished_at INTEGER
  69. );
  70. CREATE TABLE IF NOT EXISTS creation_search_jobs (
  71. id INTEGER PRIMARY KEY AUTOINCREMENT,
  72. run_id TEXT NOT NULL,
  73. query TEXT NOT NULL,
  74. platform TEXT NOT NULL,
  75. status TEXT NOT NULL,
  76. attempts INTEGER NOT NULL DEFAULT 0,
  77. searched_count INTEGER NOT NULL DEFAULT 0,
  78. display_count INTEGER NOT NULL DEFAULT 0,
  79. error TEXT,
  80. started_at INTEGER,
  81. finished_at INTEGER,
  82. UNIQUE(run_id, query, platform)
  83. );
  84. CREATE INDEX IF NOT EXISTS ix_csj_query ON creation_search_jobs(query);
  85. CREATE INDEX IF NOT EXISTS ix_csj_run_platform ON creation_search_jobs(run_id, platform);
  86. CREATE TABLE IF NOT EXISTS creation_search_items (
  87. id INTEGER PRIMARY KEY AUTOINCREMENT,
  88. run_id TEXT NOT NULL,
  89. query TEXT NOT NULL,
  90. platform TEXT NOT NULL,
  91. rank INTEGER NOT NULL,
  92. source_id TEXT,
  93. url TEXT,
  94. title TEXT,
  95. author TEXT,
  96. body_text TEXT,
  97. cover_url TEXT,
  98. image_urls TEXT NOT NULL DEFAULT '[]',
  99. video_url TEXT,
  100. media_json TEXT NOT NULL DEFAULT '{}',
  101. raw_json TEXT NOT NULL DEFAULT '{}',
  102. is_displayable INTEGER NOT NULL DEFAULT 0,
  103. error TEXT,
  104. created_at INTEGER NOT NULL,
  105. UNIQUE(run_id, query, platform, rank)
  106. );
  107. CREATE INDEX IF NOT EXISTS ix_csi_query ON creation_search_items(query);
  108. CREATE INDEX IF NOT EXISTS ix_csi_run_platform ON creation_search_items(run_id, platform);
  109. CREATE TABLE IF NOT EXISTS creation_item_classifications (
  110. item_id INTEGER PRIMARY KEY,
  111. is_creation INTEGER,
  112. reason TEXT,
  113. knowledge TEXT,
  114. prompt_version TEXT,
  115. error TEXT,
  116. classified_at INTEGER NOT NULL,
  117. FOREIGN KEY(item_id) REFERENCES creation_search_items(id) ON DELETE CASCADE
  118. );
  119. """
  120. # 旧库迁移:post_class 早期没有 knowledge/points 列,补上(已存在则忽略)
  121. _MIGRATIONS = [
  122. "ALTER TABLE post_class ADD COLUMN knowledge TEXT",
  123. "ALTER TABLE post_class ADD COLUMN points TEXT",
  124. ]
  125. def resolve_db_path(db_path: Path | str | None = None, *, env_file: str | Path | None = None) -> Path:
  126. """Resolve the SQLite path at call time.
  127. `DB_PATH` stays as the default new-data location for compatibility, while
  128. CK_SQLITE_PATH can point old demo/API runs at legacy_data/app.db.
  129. """
  130. if db_path:
  131. p = Path(db_path)
  132. else:
  133. source = load_env_file(env_file or os.getenv("CK_ENV_FILE", ".env"))
  134. raw = os.getenv("CK_SQLITE_PATH") or source.get("CK_SQLITE_PATH") or str(DB_PATH)
  135. p = Path(raw)
  136. return p if p.is_absolute() else ROOT / p
  137. def connect(db_path: Path | str | None = None) -> sqlite3.Connection:
  138. """打开(必要时建)库,建表,行按 dict 取。busy_timeout 让多进程并发写不报锁。"""
  139. p = resolve_db_path(db_path)
  140. p.parent.mkdir(parents=True, exist_ok=True)
  141. conn = sqlite3.connect(str(p), timeout=30)
  142. conn.row_factory = sqlite3.Row
  143. conn.execute("PRAGMA busy_timeout=30000") # 撞锁等 30s 而非报错(与微信补下载并发安全)
  144. conn.executescript(_SCHEMA)
  145. for sql in _MIGRATIONS:
  146. try:
  147. conn.execute(sql)
  148. except sqlite3.OperationalError:
  149. pass # 列已存在
  150. return conn
  151. def _now(ts: Optional[int]) -> int:
  152. # 调用方传入时间戳(脚本侧用 int(time.time()));缺省 0,避免本模块碰 Date/random。
  153. return int(ts) if ts is not None else 0
  154. def _record_run(conn: sqlite3.Connection, run_id: str, kind: str, note: str, ts: int) -> None:
  155. conn.execute(
  156. "INSERT OR REPLACE INTO runs(run_id, kind, note, created_at) VALUES(?,?,?,?)",
  157. (run_id, kind, note, ts),
  158. )
  159. # ---- 写入 ----------------------------------------------------------------
  160. # 哪些键是「轴」(除 query 外的结构化字段),存进 axes blob
  161. _NON_AXIS = {"query"}
  162. def insert_queries(conn: sqlite3.Connection, run_id: str, method: str,
  163. items: Iterable[dict], *, ts: Optional[int] = None) -> int:
  164. """写一批生成的 query。items=[{query, <各轴>...}]。同 (run_id,method) 先清后写(幂等)。"""
  165. t = _now(ts)
  166. conn.execute("DELETE FROM queries WHERE run_id=? AND method=?", (run_id, method))
  167. rows = []
  168. for it in items:
  169. q = (it or {}).get("query")
  170. if not q:
  171. continue
  172. axes = {k: v for k, v in it.items() if k not in _NON_AXIS}
  173. rows.append((run_id, method, q, json.dumps(axes, ensure_ascii=False), t))
  174. conn.executemany(
  175. "INSERT INTO queries(run_id, method, query, axes, created_at) VALUES(?,?,?,?,?)", rows)
  176. _record_run(conn, run_id, "queries", method, t)
  177. conn.commit()
  178. return len(rows)
  179. def _flatten_search(rec: dict) -> list[dict]:
  180. """把一条 {method, query, douyin/weixin/xiaohongshu:{ok:[...],error}} 拍成「每个结果一行」。
  181. 某渠道有结果 → 每个结果一行 ok=1;空/失败 → 一行 ok=0 记 error(保留「搜过但无结果」状态)。
  182. 只处理记录里实际出现的渠道(向后兼容只有抖音/微信的旧记录)。"""
  183. out = []
  184. method, query = rec.get("method", ""), rec.get("query", "")
  185. for platform in ("douyin", "weixin", "xiaohongshu"):
  186. if platform not in rec:
  187. continue
  188. v = rec.get(platform) or {}
  189. hits = v.get("ok") or []
  190. if hits:
  191. for h in hits:
  192. extra = {k: h[k] for k in ("nick", "images", "body_text") if k in h}
  193. out.append({"platform": platform, "ok": 1, "title": h.get("title"),
  194. "url": h.get("url"), "cover": h.get("cover"),
  195. "video": h.get("video"), "extra": extra})
  196. else:
  197. out.append({"platform": platform, "ok": 0, "extra": {"error": v.get("error") or "未搜到"}})
  198. return [{**r, "method": method, "query": query} for r in out]
  199. def insert_search_results(conn: sqlite3.Connection, run_id: str,
  200. records: Iterable[dict], *, ts: Optional[int] = None) -> int:
  201. """写一批真实搜索结果。records 为 run_search.py 产出的嵌套结构,按平台拍平入库。"""
  202. t = _now(ts)
  203. conn.execute("DELETE FROM search_results WHERE run_id=?", (run_id,))
  204. rows = []
  205. for rec in records:
  206. for r in _flatten_search(rec):
  207. rows.append((run_id, r["method"], r["query"], r["platform"], r["ok"],
  208. r.get("title"), r.get("url"), r.get("cover"), r.get("video"),
  209. json.dumps(r.get("extra", {}), ensure_ascii=False), t))
  210. conn.executemany(
  211. "INSERT INTO search_results(run_id, method, query, platform, ok, title, url, "
  212. "cover, video, extra, created_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)", rows)
  213. _record_run(conn, run_id, "search", f"{len(rows)} rows", t)
  214. conn.commit()
  215. return len(rows)
  216. # ---- 读取(分页 + 筛选)---------------------------------------------------
  217. def _page(conn: sqlite3.Connection, base: str, where: list[str], params: list[Any],
  218. page: int, size: int) -> dict:
  219. """通用分页:返回 {total, page, size, items}。"""
  220. size = max(1, min(int(size), 200))
  221. page = max(1, int(page))
  222. clause = (" WHERE " + " AND ".join(where)) if where else ""
  223. total = conn.execute(f"SELECT COUNT(*) FROM ({base}{clause})", params).fetchone()[0]
  224. rows = conn.execute(f"{base}{clause} ORDER BY id LIMIT ? OFFSET ?",
  225. params + [size, (page - 1) * size]).fetchall()
  226. return {"total": total, "page": page, "size": size, "items": [dict(r) for r in rows]}
  227. def list_queries(conn: sqlite3.Connection, *, method: Optional[str] = None,
  228. run_id: Optional[str] = None, page: int = 1, size: int = 30) -> dict:
  229. where, params = [], []
  230. if method:
  231. where.append("method=?"); params.append(method)
  232. if run_id:
  233. where.append("run_id=?"); params.append(run_id)
  234. res = _page(conn, "SELECT * FROM queries", where, params, page, size)
  235. for it in res["items"]:
  236. it["axes"] = json.loads(it.get("axes") or "{}")
  237. return res
  238. def list_search(conn: sqlite3.Connection, *, platform: Optional[str] = None,
  239. method: Optional[str] = None, ok: Optional[bool] = None,
  240. run_id: Optional[str] = None, query: Optional[str] = None,
  241. page: int = 1, size: int = 30) -> dict:
  242. where, params = [], []
  243. if platform:
  244. where.append("platform=?"); params.append(platform)
  245. if method:
  246. where.append("method=?"); params.append(method)
  247. if ok is not None:
  248. where.append("ok=?"); params.append(1 if ok else 0)
  249. if run_id:
  250. where.append("run_id=?"); params.append(run_id)
  251. if query:
  252. where.append("query=?"); params.append(query) # 按 query 文本取该条的全部搜索结果
  253. res = _page(conn, "SELECT * FROM search_results", where, params, page, size)
  254. cm = class_map(conn, {it.get("url") for it in res["items"] if it.get("url")})
  255. for it in res["items"]:
  256. it["extra"] = json.loads(it.get("extra") or "{}")
  257. it["cls"] = cm.get(it.get("url")) # 创作知识分类(无则 None=未分类)
  258. return res
  259. # ---- 帖子级创作知识分类 ----------------------------------------------------
  260. def get_unclassified_posts(conn: sqlite3.Connection) -> list[dict]:
  261. """库里有结果、但还没分类的去重帖子(按 url)。带 title/body_text/本地图片,供 classify 用。"""
  262. rows = conn.execute(
  263. "SELECT url, platform, title, cover, extra FROM search_results "
  264. "WHERE ok=1 AND url IS NOT NULL AND url NOT IN (SELECT url FROM post_class) "
  265. "GROUP BY url"
  266. ).fetchall()
  267. out = []
  268. for r in rows:
  269. extra = json.loads(r["extra"] or "{}")
  270. imgs = extra.get("images") or ([r["cover"]] if r["cover"] else [])
  271. out.append({"url": r["url"], "platform": r["platform"], "title": r["title"] or "",
  272. "body_text": extra.get("body_text", ""), "images": imgs})
  273. return out
  274. def upsert_class(conn: sqlite3.Connection, url: str, is_creation: int, reason: str, ts: int,
  275. knowledge: str = "", points: str = "") -> None:
  276. conn.execute(
  277. "INSERT OR REPLACE INTO post_class(url, is_creation, reason, knowledge, points, created_at) "
  278. "VALUES(?,?,?,?,?,?)", (url, is_creation, reason, knowledge, points, ts))
  279. conn.commit()
  280. def posts_to_classify(conn: sqlite3.Connection, platforms: list[str]) -> list[dict]:
  281. """指定平台的全部去重帖(ok=1),带 title/body_text/本地图/本地视频,供(重)分类。
  282. 不管是否已分类——上层 upsert 覆盖,便于换更忠实的判法重判。"""
  283. qs = ",".join("?" * len(platforms))
  284. rows = conn.execute(
  285. f"SELECT url, platform, title, cover, video, extra FROM search_results "
  286. f"WHERE ok=1 AND url IS NOT NULL AND platform IN ({qs}) GROUP BY url", platforms
  287. ).fetchall()
  288. out = []
  289. for r in rows:
  290. e = json.loads(r["extra"] or "{}")
  291. imgs = e.get("images") or ([r["cover"]] if r["cover"] else [])
  292. out.append({"url": r["url"], "platform": r["platform"], "title": r["title"] or "",
  293. "body_text": e.get("body_text", ""), "images": imgs, "video": r["video"]})
  294. return out
  295. def class_map(conn: sqlite3.Connection, urls) -> dict:
  296. """{url: {is_creation, reason, knowledge, points}},给取数时挂分类 + 提取的知识点。"""
  297. urls = [u for u in urls if u]
  298. if not urls:
  299. return {}
  300. qs = ",".join("?" * len(urls))
  301. rows = conn.execute(
  302. f"SELECT url, is_creation, reason, knowledge, points FROM post_class WHERE url IN ({qs})", urls
  303. ).fetchall()
  304. out = {}
  305. for r in rows:
  306. out[r["url"]] = {"is_creation": r["is_creation"], "reason": r["reason"],
  307. "knowledge": r["knowledge"] or "",
  308. "points": json.loads(r["points"]) if r["points"] else []}
  309. return out
  310. def class_counts(conn: sqlite3.Connection) -> dict:
  311. y = conn.execute("SELECT COUNT(*) FROM post_class WHERE is_creation=1").fetchone()[0]
  312. n = conn.execute("SELECT COUNT(*) FROM post_class WHERE is_creation=0").fetchone()[0]
  313. return {"creation": y, "non_creation": n}
  314. # ---- 微信正文补下载 --------------------------------------------------------
  315. def weixin_urls_missing_body(conn: sqlite3.Connection) -> list[str]:
  316. """所有微信帖(distinct url, ok=1)中 extra 还没补正文(body_text)的,供补下载。"""
  317. rows = conn.execute(
  318. "SELECT url, extra FROM search_results WHERE platform='weixin' AND ok=1 "
  319. "AND url IS NOT NULL GROUP BY url"
  320. ).fetchall()
  321. return [r["url"] for r in rows if not json.loads(r["extra"] or "{}").get("body_text")]
  322. def update_post_content(conn: sqlite3.Connection, url: str, body_text: str,
  323. images: list[str]) -> None:
  324. """把正文 + 本地图片写回该 url 的所有行的 extra(其余字段保留)。"""
  325. for r in conn.execute("SELECT id, extra FROM search_results WHERE url=?", (url,)).fetchall():
  326. e = json.loads(r["extra"] or "{}")
  327. e["body_text"] = body_text
  328. if images:
  329. e["images"] = images
  330. conn.execute("UPDATE search_results SET extra=? WHERE id=?",
  331. (json.dumps(e, ensure_ascii=False), r["id"]))
  332. conn.commit()
  333. def search_summary(conn: sqlite3.Connection) -> dict:
  334. """每条 query 搜到几个结果:{query文本: {total, ok}}。给列表页判断按钮显不显示、显几个。"""
  335. rows = conn.execute(
  336. "SELECT query, COUNT(*) AS total, SUM(ok) AS ok_n FROM search_results GROUP BY query"
  337. ).fetchall()
  338. return {r["query"]: {"total": r["total"], "ok": r["ok_n"] or 0} for r in rows}
  339. def list_runs(conn: sqlite3.Connection) -> list[dict]:
  340. rows = conn.execute("SELECT * FROM runs ORDER BY created_at DESC, run_id").fetchall()
  341. return [dict(r) for r in rows]
  342. def list_methods(conn: sqlite3.Connection, table: str = "queries") -> list[str]:
  343. """某表里有哪些 method(给前端筛选下拉)。"""
  344. if table not in ("queries", "search_results"):
  345. raise ValueError(table)
  346. rows = conn.execute(f"SELECT DISTINCT method FROM {table} ORDER BY method").fetchall()
  347. return [r[0] for r in rows]
  348. # ---- 430 query 真实采集 ----------------------------------------------------
  349. CREATION_PLATFORMS = ("xiaohongshu", "weixin", "douyin")
  350. def create_creation_run(conn: sqlite3.Connection, run_id: str, *,
  351. total_queries: int, note: str = "", ts: Optional[int] = None,
  352. status: str = "running") -> None:
  353. t = _now(ts)
  354. conn.execute(
  355. "INSERT OR REPLACE INTO creation_search_runs"
  356. "(run_id, status, total_queries, note, started_at, finished_at) VALUES(?,?,?,?,?,NULL)",
  357. (run_id, status, int(total_queries), note, t),
  358. )
  359. conn.commit()
  360. def finish_creation_run(conn: sqlite3.Connection, run_id: str, *,
  361. status: str = "finished", ts: Optional[int] = None) -> None:
  362. conn.execute(
  363. "UPDATE creation_search_runs SET status=?, finished_at=? WHERE run_id=?",
  364. (status, _now(ts), run_id),
  365. )
  366. conn.commit()
  367. def creation_run_exists(conn: sqlite3.Connection, run_id: str) -> bool:
  368. row = conn.execute(
  369. "SELECT 1 FROM creation_search_runs WHERE run_id=? LIMIT 1", (run_id,)
  370. ).fetchone()
  371. return row is not None
  372. def ensure_creation_job(conn: sqlite3.Connection, run_id: str, query: str, platform: str,
  373. *, ts: Optional[int] = None) -> int:
  374. t = _now(ts)
  375. conn.execute(
  376. "INSERT OR IGNORE INTO creation_search_jobs"
  377. "(run_id, query, platform, status, attempts, started_at) VALUES(?,?,?,?,?,?)",
  378. (run_id, query, platform, "pending", 0, t),
  379. )
  380. conn.commit()
  381. row = conn.execute(
  382. "SELECT id FROM creation_search_jobs WHERE run_id=? AND query=? AND platform=?",
  383. (run_id, query, platform),
  384. ).fetchone()
  385. return int(row["id"])
  386. def update_creation_job(conn: sqlite3.Connection, run_id: str, query: str, platform: str,
  387. *, status: str, attempts: Optional[int] = None,
  388. searched_count: Optional[int] = None,
  389. display_count: Optional[int] = None,
  390. error: Optional[str] = None, ts: Optional[int] = None) -> None:
  391. ensure_creation_job(conn, run_id, query, platform, ts=ts)
  392. sets = ["status=?"]
  393. params: list[Any] = [status]
  394. if attempts is not None:
  395. sets.append("attempts=?"); params.append(int(attempts))
  396. if searched_count is not None:
  397. sets.append("searched_count=?"); params.append(int(searched_count))
  398. if display_count is not None:
  399. sets.append("display_count=?"); params.append(int(display_count))
  400. if error is not None:
  401. sets.append("error=?"); params.append(error)
  402. if status in ("done", "failed", "partial"):
  403. sets.append("finished_at=?"); params.append(_now(ts))
  404. elif status == "running":
  405. sets.append("started_at=?"); params.append(_now(ts))
  406. params.extend([run_id, query, platform])
  407. conn.execute(
  408. f"UPDATE creation_search_jobs SET {', '.join(sets)} "
  409. "WHERE run_id=? AND query=? AND platform=?",
  410. params,
  411. )
  412. conn.commit()
  413. def creation_job_is_done(conn: sqlite3.Connection, run_id: str, query: str, platform: str,
  414. *, display_limit: int = 5) -> bool:
  415. row = conn.execute(
  416. "SELECT status, display_count FROM creation_search_jobs "
  417. "WHERE run_id=? AND query=? AND platform=?",
  418. (run_id, query, platform),
  419. ).fetchone()
  420. return bool(row and row["status"] == "done" and int(row["display_count"] or 0) >= display_limit)
  421. def upsert_creation_item(conn: sqlite3.Connection, *, run_id: str, query: str, platform: str,
  422. rank: int, source_id: str = "", url: str = "", title: str = "",
  423. author: str = "", body_text: str = "", cover_url: str = "",
  424. image_urls: Optional[list[str]] = None, video_url: str = "",
  425. media: Optional[dict] = None, raw: Optional[dict] = None,
  426. is_displayable: bool = False, error: str = "",
  427. ts: Optional[int] = None) -> int:
  428. t = _now(ts)
  429. conn.execute(
  430. "INSERT INTO creation_search_items"
  431. "(run_id, query, platform, rank, source_id, url, title, author, body_text, cover_url, "
  432. "image_urls, video_url, media_json, raw_json, is_displayable, error, created_at) "
  433. "VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?) "
  434. "ON CONFLICT(run_id, query, platform, rank) DO UPDATE SET "
  435. "source_id=excluded.source_id, url=excluded.url, title=excluded.title, "
  436. "author=excluded.author, body_text=excluded.body_text, cover_url=excluded.cover_url, "
  437. "image_urls=excluded.image_urls, video_url=excluded.video_url, media_json=excluded.media_json, "
  438. "raw_json=excluded.raw_json, is_displayable=excluded.is_displayable, "
  439. "error=excluded.error, created_at=excluded.created_at",
  440. (
  441. run_id, query, platform, int(rank), source_id, url, title, author, body_text,
  442. cover_url, json.dumps(image_urls or [], ensure_ascii=False), video_url,
  443. json.dumps(media or {}, ensure_ascii=False), json.dumps(raw or {}, ensure_ascii=False),
  444. 1 if is_displayable else 0, error, t,
  445. ),
  446. )
  447. conn.commit()
  448. row = conn.execute(
  449. "SELECT id FROM creation_search_items WHERE run_id=? AND query=? AND platform=? AND rank=?",
  450. (run_id, query, platform, int(rank)),
  451. ).fetchone()
  452. return int(row["id"])
  453. def upsert_creation_classification(conn: sqlite3.Connection, item_id: int,
  454. is_creation: Optional[int], reason: str = "",
  455. knowledge: str = "", prompt_version: str = "",
  456. error: str = "", ts: Optional[int] = None) -> None:
  457. conn.execute(
  458. "INSERT INTO creation_item_classifications"
  459. "(item_id, is_creation, reason, knowledge, prompt_version, error, classified_at) "
  460. "VALUES(?,?,?,?,?,?,?) "
  461. "ON CONFLICT(item_id) DO UPDATE SET "
  462. "is_creation=excluded.is_creation, reason=excluded.reason, knowledge=excluded.knowledge, "
  463. "prompt_version=excluded.prompt_version, error=excluded.error, classified_at=excluded.classified_at",
  464. (int(item_id), is_creation, reason, knowledge, prompt_version, error, _now(ts)),
  465. )
  466. conn.commit()
  467. def latest_creation_run_id(conn: sqlite3.Connection) -> Optional[str]:
  468. row = conn.execute(
  469. "SELECT run_id FROM creation_search_runs ORDER BY started_at DESC, run_id DESC LIMIT 1"
  470. ).fetchone()
  471. return row["run_id"] if row else None
  472. def creation_search_summary(conn: sqlite3.Connection, run_id: Optional[str] = None) -> dict:
  473. rid = run_id or latest_creation_run_id(conn)
  474. if not rid:
  475. return {"run_id": None, "queries": {}}
  476. def default_query_entry() -> dict:
  477. return {
  478. "platforms": {},
  479. "done": 0,
  480. "failed": 0,
  481. "display_count": 0,
  482. "imgtext_creation_count": 0,
  483. "imgtext_classified_count": 0,
  484. "imgtext_total_count": 0,
  485. "imgtext_target_count": 10,
  486. }
  487. rows = conn.execute(
  488. "SELECT query, platform, status, attempts, searched_count, display_count, error "
  489. "FROM creation_search_jobs WHERE run_id=? ORDER BY query, platform",
  490. (rid,),
  491. ).fetchall()
  492. out: dict[str, dict] = {}
  493. for r in rows:
  494. q = r["query"]
  495. ent = out.setdefault(q, default_query_entry())
  496. p = {
  497. "status": r["status"],
  498. "attempts": r["attempts"],
  499. "searched_count": r["searched_count"],
  500. "display_count": r["display_count"],
  501. "error": r["error"],
  502. }
  503. ent["platforms"][r["platform"]] = p
  504. ent["display_count"] += r["display_count"] or 0
  505. if r["status"] == "done":
  506. ent["done"] += 1
  507. if r["status"] == "failed":
  508. ent["failed"] += 1
  509. count_rows = conn.execute(
  510. "SELECT i.query, COUNT(*) AS total_count, "
  511. "SUM(CASE WHEN c.item_id IS NOT NULL THEN 1 ELSE 0 END) AS classified_count, "
  512. "SUM(CASE WHEN c.is_creation=1 THEN 1 ELSE 0 END) AS creation_count "
  513. "FROM creation_search_items i "
  514. "LEFT JOIN creation_item_classifications c ON c.item_id=i.id "
  515. "WHERE i.run_id=? AND i.is_displayable=1 "
  516. "AND i.platform IN ('xiaohongshu', 'weixin') "
  517. "GROUP BY i.query",
  518. (rid,),
  519. ).fetchall()
  520. for r in count_rows:
  521. ent = out.setdefault(r["query"], default_query_entry())
  522. ent["imgtext_creation_count"] = int(r["creation_count"] or 0)
  523. ent["imgtext_classified_count"] = int(r["classified_count"] or 0)
  524. ent["imgtext_total_count"] = int(r["total_count"] or 0)
  525. return {"run_id": rid, "queries": out}
  526. def get_creation_query_detail(conn: sqlite3.Connection, query: str, *,
  527. run_id: Optional[str] = None, display_limit: int = 5) -> dict:
  528. rid = run_id or latest_creation_run_id(conn)
  529. platforms = {p: {"status": "not_started", "items": [], "error": None} for p in CREATION_PLATFORMS}
  530. if not rid:
  531. return {"run_id": None, "query": query, "platforms": platforms}
  532. job_rows = conn.execute(
  533. "SELECT platform, status, attempts, searched_count, display_count, error "
  534. "FROM creation_search_jobs WHERE run_id=? AND query=?",
  535. (rid, query),
  536. ).fetchall()
  537. for r in job_rows:
  538. platforms[r["platform"]].update({
  539. "status": r["status"],
  540. "attempts": r["attempts"],
  541. "searched_count": r["searched_count"],
  542. "display_count": r["display_count"],
  543. "error": r["error"],
  544. })
  545. item_rows = conn.execute(
  546. "SELECT i.*, c.is_creation, c.reason AS class_reason, c.knowledge, "
  547. "c.prompt_version, c.error AS class_error, c.classified_at "
  548. "FROM creation_search_items i "
  549. "LEFT JOIN creation_item_classifications c ON c.item_id=i.id "
  550. "WHERE i.run_id=? AND i.query=? AND i.is_displayable=1 "
  551. "ORDER BY i.platform, i.rank",
  552. (rid, query),
  553. ).fetchall()
  554. per_platform_counts = {p: 0 for p in CREATION_PLATFORMS}
  555. for r in item_rows:
  556. p = r["platform"]
  557. if p not in platforms or per_platform_counts[p] >= display_limit:
  558. continue
  559. item = dict(r)
  560. item["image_urls"] = json.loads(item.get("image_urls") or "[]")
  561. item["media"] = json.loads(item.pop("media_json") or "{}")
  562. item.pop("raw_json", None) # 原始回包很大;详情 API 只返回前端展示需要的字段
  563. item["classification"] = {
  564. "is_creation": item.pop("is_creation"),
  565. "reason": item.pop("class_reason") or "",
  566. "knowledge": item.pop("knowledge") or "",
  567. "prompt_version": item.pop("prompt_version") or "",
  568. "error": item.pop("class_error") or "",
  569. "classified_at": item.pop("classified_at"),
  570. }
  571. platforms[p]["items"].append(item)
  572. per_platform_counts[p] += 1
  573. return {"run_id": rid, "query": query, "platforms": platforms}
  574. def creation_items_to_classify(conn: sqlite3.Connection, *, run_id: Optional[str] = None,
  575. platforms: Optional[list[str]] = None,
  576. retry_failed: bool = True,
  577. limit: Optional[int] = None) -> list[dict]:
  578. rid = run_id or latest_creation_run_id(conn)
  579. if not rid:
  580. return []
  581. where = ["i.run_id=?", "i.is_displayable=1"]
  582. params: list[Any] = [rid]
  583. if platforms:
  584. qs = ",".join("?" * len(platforms))
  585. where.append(f"i.platform IN ({qs})")
  586. params.extend(platforms)
  587. if retry_failed:
  588. where.append("(c.item_id IS NULL OR c.is_creation IS NULL)")
  589. else:
  590. where.append("c.item_id IS NULL")
  591. sql = (
  592. "SELECT i.id, i.platform, i.title, i.body_text, i.image_urls, i.video_url "
  593. "FROM creation_search_items i "
  594. "LEFT JOIN creation_item_classifications c ON c.item_id=i.id "
  595. f"WHERE {' AND '.join(where)} ORDER BY i.platform, i.id"
  596. )
  597. if limit is not None and int(limit) > 0:
  598. sql += " LIMIT ?"
  599. params.append(int(limit))
  600. rows = conn.execute(sql, params).fetchall()
  601. out = []
  602. for r in rows:
  603. d = dict(r)
  604. d["image_urls"] = json.loads(d.get("image_urls") or "[]")
  605. out.append(d)
  606. return out