Files
deep_research/scripts/runtime/sources.py
T

84 lines
3.0 KiB
Python

"""Source registry helpers for Phase 2 packets."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
def _source_key(source: dict[str, Any]) -> str:
return (source.get("id") or source.get("source_id") or source.get("doi") or source.get("url") or "").strip()
def append_packet_sources(sources_path: Path, packet: dict[str, Any]) -> int:
"""Append packet sources to sources.jsonl, preserving every citeable source_id."""
sources_path.parent.mkdir(parents=True, exist_ok=True)
existing: set[str] = set()
if sources_path.exists():
for line in sources_path.read_text(encoding="utf-8").splitlines():
if not line.strip():
continue
try:
existing.add(_source_key(json.loads(line)))
except json.JSONDecodeError:
continue
written = 0
with sources_path.open("a", encoding="utf-8") as f:
for source in packet.get("sources") or []:
key = _source_key(source)
if not key or key in existing:
continue
existing.add(key)
f.write(json.dumps(source, ensure_ascii=False) + "\n")
written += 1
return written
def rebuild_sources_from_packets(project_root: Path) -> int:
"""Rebuild phase2/sources.jsonl from packet-level source metadata.
The registry is keyed by source_id, not URL. Two packet sources may point to
the same URL but have different source_ids already cited in drafts; dropping
either row would break citation traceability.
"""
packets_dir = project_root / "phase2" / "packets"
sources_path = project_root / "phase2" / "sources.jsonl"
sources_path.parent.mkdir(parents=True, exist_ok=True)
existing_by_key: dict[str, dict[str, Any]] = {}
if sources_path.exists():
for line in sources_path.read_text(encoding="utf-8").splitlines():
if not line.strip():
continue
try:
row = json.loads(line)
except json.JSONDecodeError:
continue
key = _source_key(row)
if key:
existing_by_key[key] = row
seen: set[str] = set()
rows: list[dict[str, Any]] = []
for packet_path in sorted(packets_dir.glob("*.json")):
try:
packet = json.loads(packet_path.read_text(encoding="utf-8"))
except json.JSONDecodeError:
continue
for source in packet.get("sources") or []:
if not isinstance(source, dict):
continue
key = _source_key(source)
if not key or key in seen:
continue
seen.add(key)
previous = existing_by_key.get(key, {})
rows.append({**source, **{k: v for k, v in previous.items() if k.startswith("cache") or k.startswith("cached_")}})
sources_path.write_text(
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
encoding="utf-8",
)
return len(rows)