sync/run_sync.py

"""Weekly delta-sync entrypoint: embed new/changed papers and publish the snapshot to HF Hub."""
import json
import os
from datetime import datetime, timezone

from common.db import init_db, upsert_papers, mark_inactive_ids, update_metadata_fields
from common.vector_index import create_index, load_index, save_index, add_vector
from sync.anthology_source import iter_papers
from sync.delta import compute_delta, build_state
from sync.embeddings_client import EmbeddingsClient
from sync.hf_dataset_store import download_snapshot, upload_snapshot


def _embedding_text(paper: dict) -> str:
    if paper["abstract"]:
        return f"{paper['title']}\n\n{paper['abstract']}"
    return paper["title"]


def run_sync(anthology_path: str, hf_repo_id: str, hf_token: str,
             embedding_base_url: str, embedding_api_key: str, work_dir: str) -> None:
    os.makedirs(work_dir, exist_ok=True)
    downloaded_state = download_snapshot(hf_repo_id, work_dir, hf_token)
    last_state = downloaded_state.get("papers", {})

    db_path = os.path.join(work_dir, "papers.db")
    index_path = os.path.join(work_dir, "index.faiss")
    conn = init_db(db_path)
    try:
        index = load_index(index_path)
    except FileNotFoundError:
        index = None  # created lazily once the embedding dimension is known

    current_papers = list(iter_papers(anthology_path))
    delta = compute_delta(current_papers, last_state)

    if delta.new_or_changed:
        client = EmbeddingsClient(base_url=embedding_base_url, api_key=embedding_api_key)
        texts = [_embedding_text(p) for p in delta.new_or_changed]
        vectors = client.embed_batch(texts)  # raises RuntimeError on failure, propagates before any writes below
        # qwen3-embedding-4b returns a fixed native dim (no matryoshka), so the
        # index dimension is whatever the model emits, derived from the first vector.
        embedding_dim = len(vectors[0])
        if index is None:
            index = create_index(dim=embedding_dim)
        elif index.d != embedding_dim:
            raise RuntimeError(
                f"Embedding dimension {embedding_dim} does not match existing index "
                f"dimension {index.d}; delete the stale index.faiss (locally and on "
                f"the HF Hub) and run a fresh full sync."
            )
        upsert_rows = []
        for paper, vector in zip(delta.new_or_changed, vectors):
            faiss_id = add_vector(index, vector)
            upsert_rows.append({**paper, "active": True, "faiss_id": faiss_id})
        upsert_papers(conn, upsert_rows)

    mark_inactive_ids(conn, list(delta.removed_ids))

    # Keep bibtex/pdf_url current for all current papers without re-embedding.
    # The embedding delta above only touched new/changed papers; this
    # fills/refreshes these fields for the unchanged majority (incl. the
    # one-time backfill after the columns were added).
    update_metadata_fields(conn, current_papers)

    if index is not None:
        save_index(index, index_path)
    state = {
        "papers": build_state(current_papers),
        "last_synced_at": datetime.now(timezone.utc).isoformat(),
    }
    with open(os.path.join(work_dir, "state.json"), "w") as f:
        json.dump(state, f)

    upload_snapshot(hf_repo_id, work_dir, hf_token)


if __name__ == "__main__":
    run_sync(
        anthology_path=os.environ["ANTHOLOGY_PATH"],
        hf_repo_id=os.environ["HF_REPO_ID"],
        hf_token=os.environ["HF_TOKEN"],
        embedding_base_url=os.environ["EMBEDDING_BASE_URL"],
        embedding_api_key=os.environ["EMBEDDING_API_KEY"],
        work_dir=os.environ.get("WORK_DIR", "./work"),
    )