Upload summaries.py

metadata_search/summaries.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+from local_rag.config import LocalRagConfig
+
+TITLE_PATTERN = re.compile(r"\*\*Title:\*\*\s*(.+)", re.MULTILINE)
+_FILENAME_TITLE = re.compile(r"^.+?\s-\s\d{4}\s-\s(.+)$")
+_FILENAME_YEAR = re.compile(r"^.+?\s-\s(\d{4})\s-\s")
+
+HF_DATASET_REPO = os.getenv("WATERSCOPE_DATASET_REPO", "EnvironmentalAI/WaterScopeAI-Data")
+HF_SUMMARY_PREFIX = os.getenv("WATERSCOPE_SUMMARY_PREFIX", "Summary")
+HF_SUMMARY_CACHE = Path(
+    os.getenv("WATERSCOPE_SUMMARY_CACHE", "/tmp/waterscope-summary-cache")
+).expanduser().resolve()
+HF_SUMMARY_ON_DEMAND = os.getenv("WATERSCOPE_SUMMARY_ON_DEMAND", "1").lower() in {
+    "1", "true", "yes"
+}
+
+
+def _year_from_filename(file_name: str) -> str | None:
+    match = _FILENAME_YEAR.match(Path(file_name).name)
+    return match.group(1) if match else None
+
+
+def _download_cloud_summary(file_name: str) -> Path | None:
+    """Download one summary from the HF Dataset and cache it locally."""
+    if not HF_SUMMARY_ON_DEMAND:
+        return None
+
+    year = _year_from_filename(file_name)
+    if not year:
+        return None
+
+    repo_file = f"{HF_SUMMARY_PREFIX}/{year}/{Path(file_name).name}"
+    try:
+        downloaded = hf_hub_download(
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset",
+            filename=repo_file,
+            local_dir=str(HF_SUMMARY_CACHE),
+        )
+    except Exception:
+        return None
+
+    path = Path(downloaded)
+    return path if path.is_file() else None
+
+
+def summary_search_dirs() -> list[Path]:
+    """Directories to search for summary .txt files (first match wins)."""
+    dirs: list[Path] = []
+    for raw in (os.getenv("PAPER_SUMMARIES_DIRS") or "").split(":"):
+        raw = raw.strip()
+        if raw:
+            dirs.append(Path(raw).expanduser().resolve())
+
+    override = os.getenv("PAPER_SUMMARIES_DIR")
+    if override:
+        dirs.append(Path(override).expanduser().resolve())
+
+    project_root = Path(__file__).resolve().parents[1]
+    dirs.append(project_root / "paper_summaries")
+
+    unique: list[Path] = []
+    seen: set[str] = set()
+    for path in dirs:
+        key = str(path)
+        if key in seen:
+            continue
+        seen.add(key)
+        unique.append(path)
+    return unique
+
+
+def summaries_dir() -> Path:
+    for path in summary_search_dirs():
+        if path.exists():
+            return path
+    return Path(__file__).resolve().parents[1] / "paper_summaries"
+
+
+def _manifest_path(config: LocalRagConfig | None = None) -> Path:
+    cfg = config or LocalRagConfig.from_env()
+    return cfg.manifest_path
+
+
+def load_manifest(config: LocalRagConfig | None = None) -> dict:
+    path = _manifest_path(config)
+    if not path.exists():
+        raise FileNotFoundError(f"Manifest not found: {path}. Rebuild local_rag_data/manifest.json first.")
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def title_from_filename(file_name: str) -> str:
+    stem = Path(file_name).stem
+    match = _FILENAME_TITLE.match(stem)
+    if match:
+        return match.group(1).strip()
+    return stem
+
+
+def extract_title(file_name: str, summary: str | None = None) -> str:
+    if summary:
+        match = TITLE_PATTERN.search(summary)
+        if match:
+            return match.group(1).strip()
+    return title_from_filename(file_name)
+
+
+def summary_path(file_name: str) -> Path | None:
+    # Preserve existing local behavior first.
+    for root in summary_search_dirs():
+        candidate = root / file_name
+        if candidate.exists():
+            return candidate
+
+    # Cloud production fallback: fetch only the requested summary.
+    return _download_cloud_summary(file_name)
+
+
+def load_summary(file_name: str) -> str | None:
+    path = summary_path(file_name)
+    if path is None:
+        return None
+    return path.read_text(encoding="utf-8")
+
+
+def list_papers(config: LocalRagConfig | None = None) -> list[dict]:
+    """List indexed papers without reading every summary file into memory."""
+    manifest = load_manifest(config)
+    papers: list[dict] = []
+    for entry in manifest.get("files", []):
+        file_name = entry["file_name"]
+        papers.append(
+            {
+                "document_id": entry["document_id"],
+                "file_name": file_name,
+                "title": title_from_filename(file_name),
+                "chunk_count": entry.get("chunk_count", 0),
+                "has_summary": summary_path(file_name) is not None,
+            }
+        )
+    papers.sort(key=lambda item: item["title"].lower())
+    return papers
+
+
+def summaries_for_documents(document_ids: list[str], config: LocalRagConfig | None = None) -> list[dict]:
+    manifest = load_manifest(config)
+    by_id = {entry["document_id"]: entry for entry in manifest.get("files", [])}
+    selected: list[dict] = []
+    for document_id in document_ids:
+        entry = by_id.get(document_id)
+        if not entry:
+            continue
+        file_name = entry["file_name"]
+        summary = load_summary(file_name)
+        if not summary:
+            continue
+        selected.append(
+            {
+                "document_id": document_id,
+                "title": extract_title(file_name, summary),
+                "file_name": file_name,
+                "text": summary,
+            }
+        )
+    return selected