394 lines
15 KiB
Python
394 lines
15 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) -> 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:
|
||
report["ledger"]["muse_knowledge_draft"] = _copy_drafts(pg, sqlite)
|
||
report["ledger"]["example_knowledge_embedding"] = _copy_embeddings(pg, sqlite)
|
||
report["ledger"]["muse_knowledge_entity"] = _copy_entities(pg, sqlite)
|
||
report["ledger"]["example_ai_flavor_case"] = _copy_simple_cards(
|
||
pg, sqlite, "example_ai_flavor_case", "ai_flavor_case", "card_id", "excerpt"
|
||
)
|
||
report["ledger"]["example_ai_flavor_rule"] = _copy_simple_cards(
|
||
pg, sqlite, "example_ai_flavor_rule", "ai_flavor_rule", "name", "fix_hint"
|
||
)
|
||
report["ledger"]["example_voice_baseline"] = _copy_simple_cards(
|
||
pg, sqlite, "example_voice_baseline", "voice_baseline", "work_ref", "note"
|
||
)
|
||
report["ledger"]["muse_knowledge_document"] = _count(pg, "muse_knowledge_document")
|
||
report["ledger"]["muse_knowledge_base"] = _count(pg, "muse_knowledge_base")
|
||
report["ledger"]["example_reference_work"] = _count(pg, "example_reference_work")
|
||
report["cosine_samples"] = _cosine_samples(pg, sqlite, n=5)
|
||
report["chapters"] = _export_chapters(pg, sources_dir)
|
||
report["ledger"]["muse_content_chapter"] = report["chapters"]["pg_count"]
|
||
report["ledger"]["muse_content_block"] = _count(pg, "muse_content_block")
|
||
_copy_receipts(pg, sqlite)
|
||
report["ledger"]["example_run_receipt"] = _count(pg, "example_run_receipt")
|
||
report["ledger"]["example_candidate"] = _count(pg, "example_candidate")
|
||
report["ledger"]["example_candidate_cas"] = _count(pg, "example_candidate_cas")
|
||
report["ledger"]["example_planning_section"] = _count(pg, "example_planning_section")
|
||
report["ledger"]["example_context_freeze"] = _count(pg, "example_context_freeze")
|
||
report["reviews"] = _backfill_reviews(pg, sqlite)
|
||
report["ledger"]["example_user_decision"] = _count(pg, "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 _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:
|
||
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
|
||
card_id = f"draft:{row[1]}" if row[1] is not None else f"entity:{row[2]}"
|
||
sqlite.execute(
|
||
"""UPDATE cards SET embedding=?, content_hash=? WHERE id=?""",
|
||
(blob, row[3], card_id),
|
||
)
|
||
if sqlite.execute("SELECT changes()").fetchone()[0] == 0 and blob is not None:
|
||
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""",
|
||
(f"embedding:{row[0]}", card_id, blob, row[3]),
|
||
)
|
||
n += 1
|
||
if n % 1000 == 0:
|
||
sqlite.commit()
|
||
sqlite.commit()
|
||
if n != pg_count:
|
||
raise RuntimeError(f"example_knowledge_embedding 对账失败 pg={pg_count} sqlite={n}")
|
||
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])
|
||
card_id = f"draft:{row[1]}" if row[1] is not None else f"entity:{row[2]}"
|
||
blob_row = sqlite.execute(
|
||
"SELECT embedding FROM cards WHERE id=? AND embedding IS NOT NULL",
|
||
(card_id,),
|
||
).fetchone()
|
||
if blob_row is None:
|
||
blob_row = sqlite.execute(
|
||
"SELECT embedding FROM cards WHERE id=?",
|
||
(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) -> None:
|
||
cur = pg.execute(
|
||
"SELECT run_id, adapter_role, stage_kind, requested_model_id, actual_model_id, usage, safe_summary, create_time FROM example_run_receipt"
|
||
)
|
||
for row in cur:
|
||
run_id = row[0] or f"receipt-{row[7]}"
|
||
sqlite.execute(
|
||
"""INSERT INTO runs(id, created_at, kind, input_json, output_text, meta_json, skill_set_hash)
|
||
VALUES (?, ?, ?, ?, '', ?, 'migrated')
|
||
ON CONFLICT(id) DO NOTHING""",
|
||
(
|
||
run_id,
|
||
str(row[7] or ""),
|
||
row[1] or row[2] or "receipt",
|
||
_json({"requested_model_id": row[3], "actual_model_id": row[4], "usage": row[5]}),
|
||
_json(row[6] or {}),
|
||
),
|
||
)
|
||
sqlite.commit()
|
||
|
||
|
||
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():
|
||
print(f"COUNT {name}={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"))
|
||
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))
|
||
if report["reviews"] != 6:
|
||
return 1
|
||
if report["sqlite_bytes"] <= 0:
|
||
return 1
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|