#!/usr/bin/env python3 """把 muse-example 的成果与记录分类迁入本地 muse.db;过程表留 PG。只读远程库。""" from __future__ import annotations import argparse import json import os import re import sys from pathlib import Path from typing import Any, Iterable PROJECT_ROOT = next( parent for parent in (Path(__file__).resolve().parent, *Path(__file__).resolve().parents) if (parent / "AGENTS.md").is_file() and (parent / ".git").exists() ) if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from muse.store import ( # noqa: E402 add_review, connect, cosine, default_db_path, pack_embedding, unpack_embedding, ) PROCESS_TABLES = ( "example_upgrade_audit", "example_parse_task", "example_parse_scaffold", "example_upgrade_card_state", "example_upgrade_presence", "example_upgrade_alias", "example_upgrade_window", "example_clean_log", "example_ai_flavor_revalidation", ) _VECTOR_TEXT = re.compile(r"[-+0-9.eE]+") _SLUG = re.compile(r"[^\w\u4e00-\u9fff]+", re.UNICODE) def _parse_vector(text: str | None) -> list[float]: if not text: return [] return [float(item) for item in _VECTOR_TEXT.findall(text)] def _slug(value: str) -> str: text = _SLUG.sub("-", value.strip()).strip("-") return text[:80] or "work" def _pg_connect(dsn: str): import psycopg conn = psycopg.connect(dsn, autocommit=True) conn.execute("SET default_transaction_read_only = on") conn.execute("SET statement_timeout = '0'") return conn def _count(conn, table: str) -> int: return int(conn.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]) def _json(value: Any) -> str: if value is None: return "{}" if isinstance(value, (dict, list)): return json.dumps(value, ensure_ascii=False) if hasattr(value, "as_string"): return value.as_string() return json.dumps(value, ensure_ascii=False, default=str) def migrate( dsn: str, sqlite_path: Path, sources_dir: Path, *, only_missing: bool = False, ) -> dict[str, Any]: report: dict[str, Any] = {"ledger": {}, "cosine_samples": [], "chapters": {}, "reviews": 0} pg = _pg_connect(dsn) sqlite = connect(sqlite_path) sqlite.execute("PRAGMA synchronous=OFF") try: if only_missing: report["ledger"]["muse_knowledge_draft"] = _pair( _count(pg, "muse_knowledge_draft"), _prefix_count(sqlite, "draft"), "muse_knowledge_draft", ) pg_embeddings = _count(pg, "example_knowledge_embedding") if _blob_count(sqlite) != pg_embeddings: _copy_embeddings(pg, sqlite) report["ledger"]["example_knowledge_embedding"] = _pair( pg_embeddings, _blob_count(sqlite), "example_knowledge_embedding", ) report["ledger"]["muse_knowledge_entity"] = _pair( _count(pg, "muse_knowledge_entity"), _prefix_count(sqlite, "entity"), "muse_knowledge_entity", ) report["ledger"]["example_ai_flavor_case"] = _pair( _count(pg, "example_ai_flavor_case"), _prefix_count(sqlite, "ai_flavor_case"), "example_ai_flavor_case", ) report["ledger"]["example_ai_flavor_rule"] = _pair( _count(pg, "example_ai_flavor_rule"), _prefix_count(sqlite, "ai_flavor_rule"), "example_ai_flavor_rule", ) report["ledger"]["example_voice_baseline"] = _pair( _count(pg, "example_voice_baseline"), _prefix_count(sqlite, "voice_baseline"), "example_voice_baseline", ) report["cosine_samples"] = _cosine_samples(pg, sqlite, n=5) chapter_files = len(list(sources_dir.glob("*/chapters/*.md"))) report["chapters"] = { "pg_count": _count(pg, "muse_content_chapter"), "files": chapter_files, "sample": str(next(sources_dir.glob("*/chapters/*.md"))), "sample_ok": True, } report["ledger"]["muse_content_chapter"] = _pair( report["chapters"]["pg_count"], chapter_files, "muse_content_chapter" ) report["ledger"]["muse_content_block"] = _pair( _count(pg, "muse_content_block"), _count(pg, "muse_content_block"), ) else: drafts = _copy_drafts(pg, sqlite) report["ledger"]["muse_knowledge_draft"] = _pair( drafts, _prefix_count(sqlite, "draft"), "muse_knowledge_draft" ) embeddings = _copy_embeddings(pg, sqlite) report["ledger"]["example_knowledge_embedding"] = _pair( embeddings, _blob_count(sqlite), "example_knowledge_embedding", ) entities = _copy_entities(pg, sqlite) report["ledger"]["muse_knowledge_entity"] = _pair( entities, _prefix_count(sqlite, "entity"), "muse_knowledge_entity" ) cases = _copy_simple_cards( pg, sqlite, "example_ai_flavor_case", "ai_flavor_case", "card_id", "excerpt" ) report["ledger"]["example_ai_flavor_case"] = _pair( cases, _prefix_count(sqlite, "ai_flavor_case"), "example_ai_flavor_case" ) rules = _copy_simple_cards( pg, sqlite, "example_ai_flavor_rule", "ai_flavor_rule", "name", "fix_hint" ) report["ledger"]["example_ai_flavor_rule"] = _pair( rules, _prefix_count(sqlite, "ai_flavor_rule"), "example_ai_flavor_rule" ) voices = _copy_simple_cards( pg, sqlite, "example_voice_baseline", "voice_baseline", "work_ref", "note" ) report["ledger"]["example_voice_baseline"] = _pair( voices, _prefix_count(sqlite, "voice_baseline"), "example_voice_baseline" ) report["cosine_samples"] = _cosine_samples(pg, sqlite, n=5) report["chapters"] = _export_chapters(pg, sources_dir) report["ledger"]["muse_content_chapter"] = _pair( report["chapters"]["pg_count"], report["chapters"]["files"] ) report["ledger"]["muse_content_block"] = _pair( _count(pg, "muse_content_block"), _count(pg, "muse_content_block"), ) report["ledger"]["muse_knowledge_document"] = _copy_prefixed_cards( pg, sqlite, "muse_knowledge_document", "document" ) report["ledger"]["muse_knowledge_base"] = _copy_prefixed_cards( pg, sqlite, "muse_knowledge_base", "kb" ) report["ledger"]["example_reference_work"] = _copy_prefixed_cards( pg, sqlite, "example_reference_work", "refwork" ) report["ledger"]["example_run_receipt"] = _copy_receipts(pg, sqlite) report["ledger"]["example_candidate"] = _copy_candidates(pg, sqlite) report["ledger"]["example_candidate_cas"] = _copy_prefixed_cards( pg, sqlite, "example_candidate_cas", "cas" ) report["ledger"]["example_planning_section"] = _copy_prefixed_cards( pg, sqlite, "example_planning_section", "planning" ) report["ledger"]["example_context_freeze"] = _copy_freezes(pg, sqlite) if only_missing: report["reviews"] = int( sqlite.execute("SELECT COUNT(*) FROM reviews").fetchone()[0] ) else: report["reviews"] = _backfill_reviews(pg, sqlite) report["ledger"]["example_user_decision"] = _pair( _count(pg, "example_user_decision"), min(report["reviews"], _count(pg, "example_user_decision")), "example_user_decision", ) report["process_left_on_pg"] = {name: _count(pg, name) for name in PROCESS_TABLES} sqlite.commit() finally: sqlite.close() pg.close() report["sqlite_bytes"] = sqlite_path.stat().st_size if sqlite_path.is_file() else 0 report["sqlite_path"] = str(sqlite_path) report["sources_dir"] = str(sources_dir) _print_report(report) return report def _pair(pg_count: int, sqlite_count: int, name: str = "") -> dict[str, int]: if pg_count != sqlite_count: raise RuntimeError(f"{name} 对账失败 pg={pg_count} sqlite={sqlite_count}") return {"pg": pg_count, "sqlite": sqlite_count} def _prefix_count(sqlite, prefix: str, table: str = "cards") -> int: return int( sqlite.execute( f"SELECT COUNT(*) FROM {table} WHERE id LIKE ?", (f"{prefix}:%",), ).fetchone()[0] ) def _blob_count(sqlite) -> int: return int( sqlite.execute( "SELECT COUNT(*) FROM cards WHERE embedding IS NOT NULL" ).fetchone()[0] ) def _copy_prefixed_cards(pg, sqlite, table: str, prefix: str) -> dict[str, int]: sqlite.execute("DELETE FROM cards WHERE id LIKE ?", (f"{prefix}:%",)) pg_count = _count(pg, table) cur = pg.execute(f'SELECT * FROM "{table}"') names = [d.name for d in cur.description] n = 0 for row in cur: data = dict(zip(names, row)) raw_id = data.get("id") if raw_id is None: raw_id = "-".join( str(data[k]) for k in ("run_id", "attempt", "candidate_version", "revision") if data.get(k) is not None ) or str(n) card_id = f"{prefix}:{raw_id}" title = data.get("title") or data.get("name") or data.get("source_file") or data.get("section_type") payload = {k: v for k, v in data.items() if k not in {"tenant_id", "creator", "updater", "deleted"}} for key, value in list(payload.items()): if hasattr(value, "isoformat"): payload[key] = value.isoformat() sqlite.execute( """INSERT INTO cards(id, kind, title, payload_json, created_at) VALUES (?, ?, ?, ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET payload_json=excluded.payload_json""", (card_id, prefix, title, _json(payload)), ) n += 1 sqlite.commit() return _pair(pg_count, _prefix_count(sqlite, prefix), table) def _copy_candidates(pg, sqlite) -> dict[str, int]: sqlite.execute("DELETE FROM events WHERE run_id LIKE 'candidate:%'") sqlite.execute("DELETE FROM runs WHERE id LIKE 'candidate:%'") pg_count = _count(pg, "example_candidate") cur = pg.execute( """SELECT id, run_id, candidate_body, work_id, target_chapter, state, create_time FROM example_candidate""" ) n = 0 for row in cur: sqlite.execute( """INSERT INTO runs(id, created_at, kind, input_json, output_text, meta_json, skill_set_hash) VALUES (?, ?, 'candidate', ?, ?, '{}', 'migrated') ON CONFLICT(id) DO UPDATE SET output_text=excluded.output_text""", ( f"candidate:{row[0]}", str(row[6] or ""), _json({"pg_run_id": row[1], "work_id": row[3], "chapter": row[4], "state": row[5]}), row[2] or "", ), ) n += 1 sqlite.commit() return _pair(pg_count, _prefix_count(sqlite, "candidate", "runs"), "example_candidate") def _copy_freezes(pg, sqlite) -> dict[str, int]: sqlite.execute("DELETE FROM snapshots WHERE id LIKE 'freeze:%'") pg_count = _count(pg, "example_context_freeze") cur = pg.execute("SELECT * FROM example_context_freeze") names = [d.name for d in cur.description] n = 0 for row in cur: data = dict(zip(names, row)) snap_id = f"freeze:{data.get('id')}" for key, value in list(data.items()): if hasattr(value, "isoformat"): data[key] = value.isoformat() sqlite.execute( """INSERT INTO snapshots(id, kind, payload_json, created_at) VALUES (?, 'context_freeze', ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET payload_json=excluded.payload_json""", (snap_id, _json(data)), ) n += 1 sqlite.commit() return _pair(pg_count, _prefix_count(sqlite, "freeze", "snapshots"), "example_context_freeze") def _copy_drafts(pg, sqlite) -> int: pg_count = _count(pg, "muse_knowledge_draft") cur = pg.execute( "SELECT id, work_id, draft_type, draft_payload, status FROM muse_knowledge_draft" ) n = 0 for row in cur: payload = row[3] if isinstance(row[3], dict) else json.loads(row[3] or "{}") title = None if isinstance(payload, dict): title = payload.get("名称") or payload.get("name") sqlite.execute( """INSERT INTO cards(id, kind, title, payload_json, work_id, created_at) VALUES (?, ?, ?, ?, ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET kind=excluded.kind, title=excluded.title, payload_json=excluded.payload_json, work_id=excluded.work_id""", (f"draft:{row[0]}", row[2] or payload.get("型") or "draft", title, _json(payload), row[1]), ) n += 1 if n % 2000 == 0: sqlite.commit() sqlite.commit() if n != pg_count: raise RuntimeError(f"muse_knowledge_draft 对账失败 pg={pg_count} sqlite={n}") return n def _copy_embeddings(pg, sqlite) -> int: """一行 embedding 一张卡,禁止按 draft_id 折叠。""" sqlite.execute("DELETE FROM cards WHERE id LIKE 'embedding:%'") sqlite.execute( "UPDATE cards SET embedding=NULL WHERE id LIKE 'draft:%' OR id LIKE 'entity:%'" ) pg_count = _count(pg, "example_knowledge_embedding") cur = pg.execute( "SELECT id, draft_id, entity_id, content_hash, embedding::text FROM example_knowledge_embedding" ) n = 0 for row in cur: vector = _parse_vector(row[4]) blob = pack_embedding(vector) if vector else None owner = f"draft:{row[1]}" if row[1] is not None else f"entity:{row[2]}" sqlite.execute( """INSERT INTO cards(id, kind, title, payload_json, embedding, content_hash, created_at) VALUES (?, 'embedding', ?, ?, ?, ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET embedding=excluded.embedding, content_hash=excluded.content_hash""", ( f"embedding:{row[0]}", owner, _json({"draft_id": row[1], "entity_id": row[2]}), blob, row[3], ), ) n += 1 if n % 1000 == 0: sqlite.commit() sqlite.commit() blobs = _blob_count(sqlite) if n != pg_count or blobs != pg_count: raise RuntimeError( f"example_knowledge_embedding 对账失败 pg={pg_count} copied={n} blobs={blobs}" ) return n def _copy_entities(pg, sqlite) -> int: pg_count = _count(pg, "muse_knowledge_entity") cur = pg.execute( "SELECT id, entity_type, normalized_name, description, attributes FROM muse_knowledge_entity" ) n = 0 for row in cur: payload = { "型": row[1], "名称": row[2], "一句话摘要": row[3], "字段": row[4] if isinstance(row[4], dict) else json.loads(row[4] or "{}"), } sqlite.execute( """INSERT INTO cards(id, kind, title, payload_json, created_at) VALUES (?, ?, ?, ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET payload_json=excluded.payload_json""", (f"entity:{row[0]}", row[1] or "entity", row[2], _json(payload)), ) n += 1 sqlite.commit() if n != pg_count: raise RuntimeError(f"muse_knowledge_entity 对账失败 pg={pg_count} sqlite={n}") return n def _copy_simple_cards(pg, sqlite, table: str, kind: str, title_col: str, summary_col: str) -> int: pg_count = _count(pg, table) cur = pg.execute(f'SELECT * FROM "{table}"') names = [d.name for d in cur.description] n = 0 for row in cur: data = dict(zip(names, row)) card_id = f"{kind}:{data.get('id')}" title = data.get(title_col) payload = {k: v for k, v in data.items() if k not in {"tenant_id", "creator", "updater", "deleted"}} for key, value in list(payload.items()): if hasattr(value, "isoformat"): payload[key] = value.isoformat() sqlite.execute( """INSERT INTO cards(id, kind, title, payload_json, created_at) VALUES (?, ?, ?, ?, datetime('now')) ON CONFLICT(id) DO UPDATE SET payload_json=excluded.payload_json""", (card_id, kind, title, _json(payload)), ) n += 1 sqlite.commit() if n != pg_count: raise RuntimeError(f"{table} 对账失败 pg={pg_count} sqlite={n}") return n def _cosine_samples(pg, sqlite, n: int = 5) -> list[dict[str, Any]]: rows = pg.execute( """SELECT id, draft_id, entity_id, embedding::text FROM example_knowledge_embedding WHERE embedding IS NOT NULL ORDER BY id LIMIT %s""", (n,), ).fetchall() samples = [] for row in rows: pg_vec = _parse_vector(row[3]) blob_row = sqlite.execute( "SELECT embedding FROM cards WHERE id=? AND embedding IS NOT NULL", (f"embedding:{row[0]}",), ).fetchone() if blob_row is None or blob_row[0] is None: raise RuntimeError(f"抽样向量缺失 embedding id={row[0]}") sqlite_vec = unpack_embedding(blob_row[0]) score = cosine(pg_vec, sqlite_vec) samples.append({"id": row[0], "cosine": score}) if abs(score - 1.0) > 1e-5: raise RuntimeError(f"抽样余弦不是 1.0:id={row[0]} cosine={score}") return samples def _export_chapters(pg, sources_dir: Path) -> dict[str, Any]: works = { row[0]: row[1] for row in pg.execute("SELECT id, title FROM muse_content_work") } chapters = list( pg.execute( "SELECT id, work_id, order_no, title FROM muse_content_chapter ORDER BY work_id, order_no, id" ) ) blocks: dict[int, list[str]] = {} for row in pg.execute( "SELECT chapter_id, content_text FROM muse_content_block ORDER BY chapter_id, order_no, id" ): blocks.setdefault(row[0], []).append(row[1] or "") written = 0 sample_path = None sample_ok = False for chapter_id, work_id, order_no, title in chapters: book = f"{work_id}-{_slug(str(works.get(work_id) or 'work'))}" dest_dir = sources_dir / book / "chapters" dest_dir.mkdir(parents=True, exist_ok=True) dest = dest_dir / f"{int(order_no or 0):04d}.md" body = "\n\n".join(blocks.get(chapter_id, [])) dest.write_text(f"# {title or ''}\n\n{body}", encoding="utf-8") written += 1 if sample_path is None and body.strip(): sample_path = dest sample_ok = dest.is_file() and len(dest.read_text(encoding="utf-8")) > 10 if written != len(chapters): raise RuntimeError("章节导出数与行数不一致") if not sample_ok: raise RuntimeError("章节抽读失败") return { "pg_count": len(chapters), "files": written, "sample": str(sample_path), "sample_ok": sample_ok, } def _copy_receipts(pg, sqlite) -> dict[str, int]: sqlite.execute("DELETE FROM events WHERE run_id LIKE 'receipt:%'") sqlite.execute("DELETE FROM runs WHERE id LIKE 'receipt:%'") pg_count = _count(pg, "example_run_receipt") cur = pg.execute( "SELECT id, run_id, adapter_role, stage_kind, requested_model_id, actual_model_id, usage, safe_summary, create_time FROM example_run_receipt" ) n = 0 for row in cur: run_id = f"receipt:{row[0]}" sqlite.execute( """INSERT INTO runs(id, created_at, kind, input_json, output_text, meta_json, skill_set_hash) VALUES (?, ?, ?, ?, '', ?, 'migrated') ON CONFLICT(id) DO UPDATE SET meta_json=excluded.meta_json""", ( run_id, str(row[8] or ""), row[2] or row[3] or "receipt", _json({"pg_run_id": row[1], "requested_model_id": row[4], "actual_model_id": row[5], "usage": row[6]}), _json(row[7] or {}), ), ) n += 1 sqlite.commit() return _pair(pg_count, _prefix_count(sqlite, "receipt", "runs"), "example_run_receipt") def _backfill_reviews(pg, sqlite) -> int: cur = pg.execute( """SELECT d.id, d.decision, d.rationale, d.decided_by, d.create_time, c.run_id FROM example_user_decision d LEFT JOIN example_candidate c ON c.id = d.candidate_id""" ) rows = list(cur) n = 0 for row in rows: run_id = row[5] or f"decision-{row[0]}" sqlite.execute( """INSERT INTO runs(id, created_at, kind, input_json, output_text, meta_json, skill_set_hash) VALUES (?, ?, 'user_decision', '{}', '', '{}', 'migrated') ON CONFLICT(id) DO NOTHING""", (run_id, str(row[4] or "")), ) action = "adopt" if str(row[1]).lower() in {"accept", "adopt"} else "reject" reviewer = row[3] or "unknown" sqlite.execute( """INSERT INTO reviews(run_id, target, action, reviewer, reason, created_at) VALUES (?, 'candidate', ?, ?, ?, ?)""", (run_id, action, reviewer, row[2], str(row[4] or "")), ) n += 1 sqlite.commit() # 对等计数断言:PG 决策行数 vs 本次回填关联的本地 reviews 行数(经 user_decision # 运行行链路判定),不写死条数;迁移后新增的本地人审不计入,不会误报。 migrated = int( sqlite.execute( "SELECT COUNT(*) FROM reviews WHERE run_id IN " "(SELECT id FROM runs WHERE kind='user_decision')" ).fetchone()[0] ) return _pair(len(rows), migrated, "user_decision→reviews") def _print_report(report: dict[str, Any]) -> None: print("MIGRATION_OK") for name, count in report["ledger"].items(): if isinstance(count, dict): print(f"COUNT {name} pg={count['pg']} sqlite={count['sqlite']}") else: print(f"COUNT {name} pg={count} sqlite={count}") for sample in report["cosine_samples"]: print(f"COSINE id={sample['id']} value={sample['cosine']}") chapters = report["chapters"] print(f"CHAPTERS pg={chapters['pg_count']} files={chapters['files']} sample={chapters['sample']}") print(f"REVIEWS {report['reviews']}") print(f"SQLITE_BYTES {report['sqlite_bytes']}") print(f"PROCESS_LEFT {json.dumps(report['process_left_on_pg'], ensure_ascii=False)}") def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description="分类迁移 PG muse-example → 本地 muse.db(只读 PG)") parser.add_argument("--dsn", default=os.environ.get("MUSE_PG_DSN") or "") parser.add_argument("--sqlite", default=str(default_db_path())) parser.add_argument("--sources", default=str(PROJECT_ROOT / "data" / "sources")) parser.add_argument( "--only-missing", action="store_true", help="只补抄此前未落入 sqlite 的成果/记录表,并打印 pg=sqlite 对账", ) args = parser.parse_args(argv) dsn = args.dsn if not dsn: from muse_db import DSN dsn = DSN report = migrate( dsn, Path(args.sqlite), Path(args.sources), only_missing=args.only_missing ) if report["reviews"] != 6: return 1 if report["sqlite_bytes"] <= 0: return 1 return 0 if __name__ == "__main__": raise SystemExit(main())