Upload paper_figures.py

metadata_search/paper_figures.py

@@ -8,6 +8,7 @@ import os
 import unicodedata
 from pathlib import Path, PurePosixPath, PureWindowsPath
 from urllib.parse import quote
+from threading import RLock
 
 from metadata_search.config import PAPER_FIGURES_DIR
 
@@ -16,7 +17,8 @@ logger = logging.getLogger(__name__)
 STEM_PREFIX_MIN = 40
 
 _index: dict[str, str] | None = None
-_manifests: dict[str, dict] | None = None
+_manifests: dict[str, dict | None] | None = None
+_index_lock = RLock()
 
 
 HF_DATASET_REPO = os.getenv(
@@ -82,15 +84,33 @@ def _batch_folders(path: Path) -> list[str]:
     if not isinstance(batch.get("papers"), list):
         raise ValueError(f"Missing papers list in media batch manifest: {path}")
     folders = []
+    seen = set()
     for record in batch["papers"]:
         folder = _media_folder(str(record.get("output_directory") or ""))
         if folder is None:
             raise ValueError("Invalid media output_directory in batch manifest")
-        if folder not in folders:
+        if folder not in seen:
             folders.append(folder)
+            seen.add(folder)
     return folders
 
 
+
+def _batch_index(path: Path) -> tuple[dict[str, str], dict[str, dict | None]]:
+    """Index folder and PDF aliases without reading individual study manifests."""
+    folders = dict.fromkeys(_batch_folders(path))
+    stem_index: dict[str, str] = {}
+    for folder in folders:
+        _register_stem(stem_index, PurePosixPath(folder).name, folder)
+    batch = json.loads(path.read_text(encoding="utf-8-sig"))
+    for record in batch["papers"]:
+        folder = _media_folder(str(record.get("output_directory") or ""))
+        source_pdf = str(record.get("source_pdf") or "").strip()
+        if source_pdf:
+            _register_stem(stem_index, PureWindowsPath(source_pdf).name, folder)
+    return stem_index, folders
+
+
 def resolve_figure_asset(folder: str, filename: str) -> Path | None:
     """Resolve one figure asset locally, then fetch only that asset if needed."""
     safe_folder = _media_folder(folder)
@@ -149,7 +169,12 @@ def _register_stem(stem_index: dict[str, str], key: str, folder: str) -> None:
     stem_index[stem] = folder
 
 
-def _load_index() -> tuple[dict[str, str], dict[str, dict]]:
+def _load_index() -> tuple[dict[str, str], dict[str, dict | None]]:
+    with _index_lock:
+        return _load_index_locked()
+
+
+def _load_index_locked() -> tuple[dict[str, str], dict[str, dict | None]]:
     global _index, _manifests
     if _index is not None and _manifests is not None:
         return _index, _manifests
@@ -157,6 +182,25 @@ def _load_index() -> tuple[dict[str, str], dict[str, dict]]:
     stem_index: dict[str, str] = {}
     manifests: dict[str, dict] = {}
 
+    # The overall manifest is sufficient for health counts and paper lookup.
+    # None marks an indexed study whose detailed manifest is not loaded yet.
+    root = figures_dir()
+    batch_path = root / "batch_manifest.json" if root is not None else None
+    if batch_path is None or not batch_path.is_file():
+        batch_path = _hf_download(f"{HF_FIGURES_PREFIX}/batch_manifest.json")
+    if batch_path is not None:
+        try:
+            _index, _manifests = _batch_index(batch_path)
+            logger.info("Indexed media for %s papers from overall manifest", len(_manifests))
+            return _index, _manifests
+        except (OSError, ValueError, TypeError, AttributeError) as exc:
+            logger.warning("Unable to read overall media manifest %s: %s", batch_path, exc)
+            # Do not permanently cache a transient failure.
+            return {}, {}
+    if root is None:
+        logger.warning("Overall media manifest unavailable; will retry on next request")
+        return {}, {}
+
     # 1) Existing local behavior.
     root = figures_dir()
     if root is not None:
@@ -182,55 +226,6 @@ def _load_index() -> tuple[dict[str, str], dict[str, dict]]:
             if source_pdf:
                 _register_stem(stem_index, PureWindowsPath(source_pdf).name, folder)
 
-    # 2) Cloud fallback. List repository paths once, then download only
-    #    manifest.json files. Image assets remain on-demand.
-    if not manifests and HF_FIGURES_ON_DEMAND:
-        try:
-            suffix = "/manifest.json"
-            prefix = f"{HF_FIGURES_PREFIX}/"
-            batch_path = _hf_download(f"{prefix}batch_manifest.json")
-            if batch_path is not None:
-                manifest_files = [f"{prefix}{folder}{suffix}"
-                                  for folder in _batch_folders(batch_path)]
-            else:
-                # Support data without an overall manifest; only study manifests.
-                repo_files = _hf_api().list_repo_files(
-                    repo_id=HF_DATASET_REPO, repo_type="dataset",
-                )
-                manifest_files = [path for path in repo_files
-                                  if path.startswith(prefix) and path.endswith(suffix)]
-
-            for repo_file in manifest_files:
-                relative = repo_file[len(prefix):]
-                folder = relative[:-len(suffix)]
-                # Keep the year component when resolving study manifests.
-                if _media_folder(folder) is None:
-                    continue
-
-                manifest_path = _hf_download(repo_file)
-                if manifest_path is None:
-                    continue
-                try:
-                    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
-                except (OSError, json.JSONDecodeError) as exc:
-                    logger.warning("Skipping cloud figure manifest %s: %s", repo_file, exc)
-                    continue
-
-                manifests[folder] = manifest
-                _register_stem(stem_index, PurePosixPath(folder).name, folder)
-
-                source_pdf = (manifest.get("source_pdf") or "").strip()
-                if source_pdf:
-                    _register_stem(stem_index, PureWindowsPath(source_pdf).name, folder)
-
-            logger.info(
-                "Indexed figures for %s papers from HF Dataset %s",
-                len(manifests),
-                HF_DATASET_REPO,
-            )
-        except Exception as exc:
-            logger.warning("Unable to build HF figure manifest index: %s", exc)
-
     _index = stem_index
     _manifests = manifests
     return _index, _manifests
@@ -299,7 +294,22 @@ def figure_image_url(folder: str, filename: str) -> str:
 
 def figures_payload_for_folder(folder: str) -> dict | None:
     _, manifests = _load_index()
-    manifest = manifests.get(folder)
+    if folder not in manifests:
+        return None
+    with _index_lock:
+        manifest = manifests.get(folder)
+        if manifest is None:
+            manifest_path = resolve_figure_asset(folder, "manifest.json")
+            if manifest_path is None:
+                return None
+            try:
+                manifest = json.loads(manifest_path.read_text(encoding="utf-8-sig"))
+                if not isinstance(manifest, dict):
+                    raise ValueError("Study manifest must be a JSON object")
+            except (OSError, ValueError) as exc:
+                logger.warning("Unable to read study manifest %s: %s", manifest_path, exc)
+                return None
+            manifests[folder] = manifest
     if not manifest:
         return None