#!/usr/bin/env python3 """Convert Deep Research source IDs into numeric citations for final output.""" from __future__ import annotations import argparse import json import re from pathlib import Path from typing import Any SRC_CITE_RE = re.compile(r"\[((?:src_[A-Za-z0-9_-]+)(?:\s*,\s*src_[A-Za-z0-9_-]+)*)\]") def load_sources(path: Path) -> dict[str, dict[str, Any]]: sources: dict[str, dict[str, Any]] = {} if not path.exists(): return sources for line in path.read_text(encoding="utf-8").splitlines(): if not line.strip(): continue try: obj = json.loads(line) except json.JSONDecodeError: continue sid = obj.get("id") or obj.get("source_id") if sid: sources[str(sid)] = obj return sources def extract_ordered_source_ids(text: str) -> list[str]: ordered: list[str] = [] seen: set[str] = set() for match in SRC_CITE_RE.finditer(text): for sid in [item.strip() for item in match.group(1).split(",")]: if sid and sid not in seen: seen.add(sid) ordered.append(sid) return ordered def _canonical_source_key(sid: str, source: dict[str, Any] | None) -> str: """Return a stable de-duplication key for a source record. Phase 2 often creates chapter-local source IDs for the same local PDF or official guideline. Final references should cite the underlying source once, while citation_map.json keeps the full src_id traceability. """ if not source: return f"missing:{sid}" title = re.sub(r"\s+", " ", str(source.get("title") or source.get("name") or sid)).strip().lower() title = title.removesuffix(" ocr").removesuffix(".ocr").strip() doi = str(source.get("doi") or "").strip().lower() if doi: return f"doi:{doi}" path = str(source.get("path") or "").strip() url = str(source.get("url") or "").strip() if title and ("phase0/extracted/" in path or "phase0/extracted/" in url): return f"local-material:{title}" for field in ("url", "path"): value = str(source.get(field) or "").strip() if value: return f"{field}:{value.rstrip('/').lower()}" return f"title:{title or sid}" def build_numeric_mapping( ordered_ids: list[str], sources: dict[str, dict[str, Any]], ) -> tuple[dict[str, int], list[dict[str, Any]]]: mapping: dict[str, int] = {} records: list[dict[str, Any]] = [] seen_keys: dict[str, int] = {} record_by_number: dict[int, dict[str, Any]] = {} for sid in ordered_ids: source = sources.get(sid) key = _canonical_source_key(sid, source) if key in seen_keys: number = seen_keys[key] mapping[sid] = number record_by_number[number].setdefault("source_ids", []).append(sid) continue number = len(records) + 1 seen_keys[key] = number mapping[sid] = number record = { "number": number, "source_id": sid, "source_ids": [sid], "source": source or {}, "dedupe_key": key, } records.append(record) record_by_number[number] = record return mapping, records def format_reference(number: int, sid: str, source: dict[str, Any] | None) -> str: if not source: return f"{number}. {sid}. （sources.jsonl 未找到该来源）" authors = ", ".join(source.get("authors", [])) if source.get("authors") else "" year = source.get("year") or source.get("date") or "" title = source.get("title") or source.get("name") or sid title = re.sub(r"(?i)(?:\s+OCR|\.ocr)$", "", str(title)).strip() publisher = source.get("publisher") or source.get("venue") or source.get("source") or "" url = source.get("url") or source.get("path") or "" parts = [f"{number}. "] if authors: parts.append(f"{authors}. ") if year: parts.append(f"({year}). ") parts.append(str(title)) if publisher: parts.append(f". {publisher}") if url: parts.append(f". {url}") return "".join(parts) def convert_citations(text: str, mapping: dict[str, int]) -> str: def repl(match: re.Match[str]) -> str: ids = [item.strip() for item in match.group(1).split(",") if item.strip()] nums: list[str] = [] seen: set[int] = set() for sid in ids: if sid not in mapping: continue number = mapping[sid] if number in seen: continue seen.add(number) nums.append(str(number)) return "[" + ", ".join(nums) + "]" if nums else match.group(0) return SRC_CITE_RE.sub(repl, text) def strip_existing_reference_section(text: str) -> str: pattern = re.compile(r"\n##\s*(?:参考文献|参考来源清单|References)\s*\n.*\Z", re.S) return pattern.sub("", text).rstrip() + "\n" def number_citations( *, text: str, sources: dict[str, dict[str, Any]], ) -> tuple[str, list[dict[str, Any]]]: ordered_ids = extract_ordered_source_ids(text) mapping, records = build_numeric_mapping(ordered_ids, sources) body = convert_citations(strip_existing_reference_section(text), mapping).rstrip() ref_lines = ["", "## 参考来源清单", ""] for record in records: ref_lines.append(format_reference(record["number"], record["source_id"], record["source"])) return body + "\n" + "\n".join(ref_lines).rstrip() + "\n", records def main() -> int: parser = argparse.ArgumentParser(description="Convert [src_xxx] citations to numeric citations") parser.add_argument("project", help="Project directory") parser.add_argument("--input", default="phase4/final_zh.md") parser.add_argument("--output", default="phase4/final_zh_numbered.md") parser.add_argument("--sources", default="phase2/sources.jsonl") parser.add_argument("--map", default="phase4/citation_map.json") args = parser.parse_args() project = Path(args.project) src_path = project / args.input out_path = project / args.output sources_path = project / args.sources map_path = project / args.map if not src_path.exists(): raise SystemExit(f"input not found: {src_path}") sources = load_sources(sources_path) numbered, records = number_citations(text=src_path.read_text(encoding="utf-8"), sources=sources) out_path.parent.mkdir(parents=True, exist_ok=True) out_path.write_text(numbered, encoding="utf-8") map_path.parent.mkdir(parents=True, exist_ok=True) map_path.write_text(json.dumps(records, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") print(f"Wrote: {out_path.relative_to(project)}") print(f"Wrote: {map_path.relative_to(project)}") print(f"Citations: {len(records)}") return 0 if __name__ == "__main__": raise SystemExit(main())