618 lines
23 KiB
Python
618 lines
23 KiB
Python
#!/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()
|
||
if n != 6:
|
||
raise RuntimeError(f"user_decision 回填 reviews 应为 6，实际 {n}")
|
||
return n
|
||
|
||
|
||
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())
|