305 lines
10 KiB
Python
305 lines
10 KiB
Python
#!/usr/bin/env python3
|
||
"""Repair ChatGPT citation tokens in Markdown exports."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import pathlib
|
||
import re
|
||
import sys
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
|
||
CITATION_RE = re.compile(r"cite(?:turn\d+(?:view|search)\d+)+")
|
||
REF_ID_RE = re.compile(r"turn\d+(?:view|search)\d+")
|
||
FENCE_RE = re.compile(r"^( {0,3})(`{3,}|~{3,})")
|
||
LINK_RE = re.compile(r"\[[^\]\n]+\]\([^) \n]+(?: [^)]+)?\)")
|
||
REFERENCE_HEADING_RE = re.compile(r"^## +参考资料\s*$", re.MULTILINE)
|
||
HTML_METADATA_RE = re.compile(
|
||
r"<!--\s*deep-research-metadata\s*(?P<json>\{.*?\})\s*-->",
|
||
re.DOTALL,
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CitationRecord:
|
||
ref_id: str
|
||
url: str
|
||
title: str = ""
|
||
attribution: str = ""
|
||
|
||
def footnote_text(self) -> str:
|
||
label = self.title or self.url
|
||
if self.attribution and self.title:
|
||
return f"{self.attribution}：{self.title}。{self.url}"
|
||
return f"{label}。{self.url}"
|
||
|
||
|
||
def extract_ref_ids(token: str) -> list[str]:
|
||
return REF_ID_RE.findall(token)
|
||
|
||
|
||
def ref_id_from_ref(ref: dict[str, Any]) -> str | None:
|
||
turn_index = ref.get("turn_index")
|
||
ref_type = ref.get("ref_type")
|
||
ref_index = ref.get("ref_index")
|
||
if turn_index is None or ref_type not in {"view", "search"} or ref_index is None:
|
||
return None
|
||
return f"turn{turn_index}{ref_type}{ref_index}"
|
||
|
||
|
||
def _first_string(*values: Any) -> str:
|
||
for value in values:
|
||
if isinstance(value, str) and value.strip():
|
||
return value.strip()
|
||
return ""
|
||
|
||
|
||
def record_from_ref(ref: dict[str, Any]) -> CitationRecord | None:
|
||
ref_id = _first_string(ref.get("ref_id"), ref.get("id")) or ref_id_from_ref(ref)
|
||
url = _first_string(ref.get("url"), ref.get("safe_url"), ref.get("source_url"))
|
||
title = _first_string(ref.get("title"), ref.get("name"), ref.get("text"))
|
||
attribution = _first_string(ref.get("attribution"), ref.get("source"), ref.get("site_name"))
|
||
if not ref_id or not url:
|
||
return None
|
||
return CitationRecord(ref_id=ref_id, url=url, title=title, attribution=attribution)
|
||
|
||
|
||
def _walk_dicts(value: Any) -> list[dict[str, Any]]:
|
||
found: list[dict[str, Any]] = []
|
||
if isinstance(value, dict):
|
||
found.append(value)
|
||
for child in value.values():
|
||
found.extend(_walk_dicts(child))
|
||
elif isinstance(value, list):
|
||
for child in value:
|
||
found.extend(_walk_dicts(child))
|
||
return found
|
||
|
||
|
||
def build_citation_map(metadata: dict[str, Any] | None) -> dict[str, CitationRecord]:
|
||
if not metadata:
|
||
return {}
|
||
|
||
citation_map: dict[str, CitationRecord] = {}
|
||
for node in _walk_dicts(metadata):
|
||
record = record_from_ref(node)
|
||
if record and record.ref_id not in citation_map:
|
||
citation_map[record.ref_id] = record
|
||
|
||
return citation_map
|
||
|
||
|
||
def split_existing_references(markdown: str) -> tuple[str, str]:
|
||
match = REFERENCE_HEADING_RE.search(markdown)
|
||
if not match:
|
||
return markdown.rstrip(), ""
|
||
return markdown[: match.start()].rstrip(), markdown[match.start() :].strip()
|
||
|
||
|
||
def collect_markdown_links(markdown: str) -> list[str]:
|
||
seen: set[str] = set()
|
||
links: list[str] = []
|
||
for link in LINK_RE.findall(markdown):
|
||
if link not in seen:
|
||
seen.add(link)
|
||
links.append(link)
|
||
return links
|
||
|
||
|
||
def replace_citations(
|
||
markdown: str, citation_map: dict[str, CitationRecord]
|
||
) -> tuple[str, list[CitationRecord], list[str]]:
|
||
footnote_numbers: dict[str, int] = {}
|
||
ordered_records: list[CitationRecord] = []
|
||
missing: list[str] = []
|
||
|
||
def marker_for_ref(ref_id: str) -> str:
|
||
if ref_id not in citation_map:
|
||
if ref_id not in missing:
|
||
missing.append(ref_id)
|
||
return ""
|
||
if ref_id not in footnote_numbers:
|
||
footnote_numbers[ref_id] = len(ordered_records) + 1
|
||
ordered_records.append(citation_map[ref_id])
|
||
return f"[^{footnote_numbers[ref_id]}]"
|
||
|
||
output: list[str] = []
|
||
in_fence = False
|
||
fence_marker = ""
|
||
|
||
for line in markdown.splitlines(keepends=True):
|
||
body = line[:-1] if line.endswith("\n") else line
|
||
newline = "\n" if line.endswith("\n") else ""
|
||
if body.endswith("\r"):
|
||
body = body[:-1]
|
||
newline = "\r" + newline
|
||
|
||
fence_match = FENCE_RE.match(body)
|
||
if fence_match:
|
||
marker = fence_match.group(2)
|
||
marker_char = marker[0]
|
||
if not in_fence:
|
||
in_fence = True
|
||
fence_marker = marker_char * len(marker)
|
||
elif marker_char == fence_marker[0] and len(marker) >= len(fence_marker):
|
||
in_fence = False
|
||
fence_marker = ""
|
||
output.append(line)
|
||
continue
|
||
|
||
if in_fence:
|
||
output.append(line)
|
||
continue
|
||
|
||
def replace_token(match: re.Match[str]) -> str:
|
||
return "".join(marker_for_ref(ref_id) for ref_id in extract_ref_ids(match.group(0)))
|
||
|
||
output.append(CITATION_RE.sub(replace_token, body) + newline)
|
||
|
||
return "".join(output), ordered_records, missing
|
||
|
||
|
||
def strip_citation_tokens(markdown: str) -> str:
|
||
lines: list[str] = []
|
||
in_fence = False
|
||
fence_marker = ""
|
||
|
||
for line in markdown.splitlines(keepends=True):
|
||
body = line[:-1] if line.endswith("\n") else line
|
||
newline = "\n" if line.endswith("\n") else ""
|
||
if body.endswith("\r"):
|
||
body = body[:-1]
|
||
newline = "\r" + newline
|
||
|
||
fence_match = FENCE_RE.match(body)
|
||
if fence_match:
|
||
marker = fence_match.group(2)
|
||
marker_char = marker[0]
|
||
if not in_fence:
|
||
in_fence = True
|
||
fence_marker = marker_char * len(marker)
|
||
elif marker_char == fence_marker[0] and len(marker) >= len(fence_marker):
|
||
in_fence = False
|
||
fence_marker = ""
|
||
lines.append(line)
|
||
continue
|
||
|
||
if in_fence:
|
||
lines.append(line)
|
||
else:
|
||
lines.append(CITATION_RE.sub("", body) + newline)
|
||
|
||
return "".join(lines)
|
||
|
||
|
||
def render_reference_section(
|
||
records: list[CitationRecord], existing_links: list[str], missing: list[str], degraded: bool
|
||
) -> str:
|
||
lines = ["## 参考资料"]
|
||
|
||
if degraded:
|
||
lines.append("")
|
||
lines.append("> 无法恢复逐句脚注映射：未找到可用元数据或部分 turn ID 缺少 URL。")
|
||
|
||
if records:
|
||
lines.append("")
|
||
for index, record in enumerate(records, start=1):
|
||
lines.append(f"[^{index}]: {record.footnote_text()}")
|
||
|
||
if existing_links:
|
||
lines.append("")
|
||
lines.append("### 原始来源链接")
|
||
for link in existing_links:
|
||
lines.append(f"- {link}")
|
||
|
||
if missing:
|
||
lines.append("")
|
||
lines.append("### 未解析引用 ID")
|
||
for ref_id in missing:
|
||
lines.append(f"- `{ref_id}`")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def repair_markdown(markdown: str, metadata: dict[str, Any] | None = None) -> str:
|
||
body, existing_reference_section = split_existing_references(markdown)
|
||
existing_links = collect_markdown_links(existing_reference_section)
|
||
citation_map = build_citation_map(metadata)
|
||
|
||
if citation_map:
|
||
replaced, records, missing = replace_citations(body, citation_map)
|
||
degraded = bool(missing)
|
||
cleaned = strip_citation_tokens(replaced) if missing else replaced
|
||
reference_section = render_reference_section(records, existing_links, missing, degraded)
|
||
return f"{cleaned.rstrip()}\n\n{reference_section}\n"
|
||
|
||
cleaned = strip_citation_tokens(body)
|
||
reference_section = render_reference_section([], existing_links, [], degraded=True)
|
||
return f"{cleaned.rstrip()}\n\n{reference_section}\n"
|
||
|
||
|
||
def load_json(path: pathlib.Path) -> dict[str, Any]:
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
value = json.load(handle)
|
||
if not isinstance(value, dict):
|
||
raise ValueError("metadata JSON must contain an object at the top level")
|
||
return value
|
||
|
||
|
||
def extract_embedded_metadata(markdown: str) -> dict[str, Any] | None:
|
||
match = HTML_METADATA_RE.search(markdown)
|
||
if not match:
|
||
return None
|
||
value = json.loads(match.group("json"))
|
||
if not isinstance(value, dict):
|
||
raise ValueError("embedded metadata JSON must contain an object at the top level")
|
||
return value
|
||
|
||
|
||
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description="Repair ChatGPT/Deep Research citation tokens in Markdown."
|
||
)
|
||
parser.add_argument("path", type=pathlib.Path, help="Markdown file to repair")
|
||
parser.add_argument("--metadata", type=pathlib.Path, help="Deep Research metadata JSON")
|
||
parser.add_argument("-o", "--output", type=pathlib.Path, help="Write repaired Markdown here")
|
||
parser.add_argument("--in-place", action="store_true", help="Rewrite the input Markdown file")
|
||
parser.add_argument("--dry-run", action="store_true", help="Print repaired Markdown to stdout")
|
||
return parser.parse_args(argv)
|
||
|
||
|
||
def main(argv: list[str]) -> int:
|
||
args = parse_args(argv)
|
||
if sum(bool(value) for value in (args.output, args.in_place, args.dry_run)) != 1:
|
||
print("error: choose exactly one of --output, --in-place, or --dry-run", file=sys.stderr)
|
||
return 2
|
||
|
||
try:
|
||
markdown = args.path.read_text(encoding="utf-8")
|
||
metadata = load_json(args.metadata) if args.metadata else extract_embedded_metadata(markdown)
|
||
fixed = repair_markdown(markdown, metadata=metadata)
|
||
except (OSError, ValueError, json.JSONDecodeError) as exc:
|
||
print(f"error: {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
if args.dry_run:
|
||
sys.stdout.write(fixed)
|
||
return 0
|
||
|
||
destination = args.path if args.in_place else args.output
|
||
assert destination is not None
|
||
try:
|
||
destination.write_text(fixed, encoding="utf-8")
|
||
except OSError as exc:
|
||
print(f"error: cannot write {destination}: {exc}", file=sys.stderr)
|
||
return 1
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main(sys.argv[1:]))
|