from __future__ import annotations import json import os import tempfile import threading import time from typing import Any, Literal, Mapping from pydantic import BaseModel, ConfigDict, Field from app.services.citation_graph import CitationGraphService CACHE_VERSION = 2 _DEFAULT_ROOT = os.path.expanduser("~/.studybuddy/citation_exploration") _RICH_METADATA_TTL = 7 * 86400 _SPARSE_METADATA_TTL = 6 * 3600 _TOPOLOGY_TTL = 30 * 86400 _ERROR_TTL = 15 * 60 class CachedReferenceSet(BaseModel): model_config = ConfigDict(extra="forbid", strict=True) provider: str work_id: str referenced_work_ids: list[str] = Field(default_factory=list) updated_at: float expires_at: float error: str = "" @property def fresh(self) -> bool: return self.expires_at > time.time() class CachedNeighborhood(BaseModel): model_config = ConfigDict(extra="forbid", strict=True) key: str exploration_kind: Literal["paper", "author"] seed_candidate_id: str = "" author_key: str = "" paper_ids: list[str] = Field(default_factory=list) primary_edges: list[dict[str, Any]] = Field(default_factory=list) induced_edges: list[dict[str, Any]] = Field(default_factory=list) primary_state: Literal["complete", "degraded"] = "complete" induced_state: Literal["pending", "complete", "degraded"] = "pending" warnings: list[str] = Field(default_factory=list) updated_at: float expires_at: float @property def fresh(self) -> bool: return self.expires_at > time.time() class ExplorationCache(BaseModel): model_config = ConfigDict(extra="forbid", strict=True) cache_version: int = CACHE_VERSION papers: dict[str, dict[str, Any]] = Field(default_factory=dict) neighborhoods: dict[str, CachedNeighborhood] = Field(default_factory=dict) reference_sets: dict[str, CachedReferenceSet] = Field(default_factory=dict) edges: dict[str, dict[str, Any]] = Field(default_factory=dict) updated_at: float = 0.0 class CitationExplorationStore: """Atomic, project-scoped cache for external paper graph exploration.""" _locks: dict[str, threading.RLock] = {} _locks_guard = threading.Lock() def __init__(self, project_id: str, root: str | None = None) -> None: if not CitationGraphService.is_valid_project_id(project_id): raise ValueError("Invalid project id.") self.project_id = project_id self.root = root or _DEFAULT_ROOT os.makedirs(self.root, exist_ok=True) with self._locks_guard: self._lock = self._locks.setdefault(self._path(), threading.RLock()) def _path(self) -> str: return os.path.join(self.root, f"{self.project_id}.json") def load(self) -> ExplorationCache: with self._lock: try: with open(self._path(), encoding="utf-8") as file: payload = json.load(file) if int(payload.get("cache_version") or 0) != CACHE_VERSION: return ExplorationCache() return ExplorationCache(**payload) except (FileNotFoundError, OSError, ValueError, TypeError, json.JSONDecodeError): return ExplorationCache() def save(self, cache: ExplorationCache) -> None: cache.updated_at = time.time() with self._lock: descriptor, temporary = tempfile.mkstemp( prefix=f".{self.project_id}.", suffix=".tmp", dir=self.root, text=True ) try: with os.fdopen(descriptor, "w", encoding="utf-8") as file: json.dump(cache.model_dump(), file, indent=2) file.flush() os.fsync(file.fileno()) os.replace(temporary, self._path()) finally: if os.path.exists(temporary): os.remove(temporary) @staticmethod def paper_is_rich(paper: Mapping[str, Any]) -> bool: authors = paper.get("author_details") or [] destinations = paper.get("destinations") or [] abstract_known = bool(paper.get("abstract")) or paper.get("abstract_status") == "unavailable" return bool(paper.get("title") and authors and destinations and abstract_known) @classmethod def paper_fresh(cls, paper: Mapping[str, Any]) -> bool: try: return float(paper.get("cache_expires_at") or 0) > time.time() except (TypeError, ValueError): return False @classmethod def prepare_paper(cls, paper: Mapping[str, Any], *, error: str = "") -> dict[str, Any]: payload = dict(paper) now = time.time() ttl = _ERROR_TTL if error else _RICH_METADATA_TTL if cls.paper_is_rich(payload) else _SPARSE_METADATA_TTL payload["cache_updated_at"] = now payload["cache_expires_at"] = now + ttl payload["metadata_quality"] = "rich" if cls.paper_is_rich(payload) else "sparse" if error: payload["last_provider_error"] = error return payload def candidate(self, candidate_id: str) -> dict[str, Any] | None: paper = self.load().papers.get(candidate_id) return dict(paper) if paper else None def upsert_papers(self, papers: list[Mapping[str, Any]]) -> ExplorationCache: cache = self.load() for incoming in papers: candidate_id = str(incoming.get("candidate_id") or "") if not candidate_id: continue prepared = self.prepare_paper(incoming, error=str(incoming.get("error") or "")) existing = cache.papers.get(candidate_id) if existing and self.paper_is_rich(existing) and not self.paper_is_rich(prepared): merged = dict(prepared) for key, value in existing.items(): if value not in (None, "", [], {}): merged[key] = value prepared = merged cache.papers[candidate_id] = prepared self.save(cache) return cache def neighborhood(self, key: str, *, fresh_only: bool = True) -> CachedNeighborhood | None: row = self.load().neighborhoods.get(key) return row if row and (row.fresh or not fresh_only) else None def save_neighborhood(self, neighborhood: CachedNeighborhood) -> ExplorationCache: cache = self.load() cache.neighborhoods[neighborhood.key] = neighborhood for edge in [*neighborhood.primary_edges, *neighborhood.induced_edges]: edge_id = str(edge.get("id") or "") if edge_id: cache.edges[edge_id] = dict(edge) self.save(cache) return cache @staticmethod def new_neighborhood( *, key: str, exploration_kind: Literal["paper", "author"], paper_ids: list[str], primary_edges: list[dict[str, Any]], seed_candidate_id: str = "", author_key: str = "", warnings: list[str] | None = None, degraded: bool = False, ) -> CachedNeighborhood: now = time.time() return CachedNeighborhood( key=key, exploration_kind=exploration_kind, seed_candidate_id=seed_candidate_id, author_key=author_key, paper_ids=paper_ids, primary_edges=primary_edges, primary_state="degraded" if degraded else "complete", warnings=warnings or [], updated_at=now, expires_at=now + (_ERROR_TTL if degraded else _TOPOLOGY_TTL), ) def reference_set(self, candidate_id: str, *, fresh_only: bool = True) -> CachedReferenceSet | None: row = self.load().reference_sets.get(candidate_id) return row if row and (row.fresh or not fresh_only) else None def save_reference_sets(self, rows: Mapping[str, CachedReferenceSet]) -> ExplorationCache: cache = self.load() cache.reference_sets.update(rows) self.save(cache) return cache @staticmethod def new_reference_set( *, provider: str, work_id: str, referenced_work_ids: list[str], error: str = "" ) -> CachedReferenceSet: now = time.time() return CachedReferenceSet( provider=provider, work_id=work_id, referenced_work_ids=sorted(set(referenced_work_ids)), updated_at=now, expires_at=now + (_ERROR_TTL if error else _TOPOLOGY_TTL), error=error, ) def merge_induced_edges( self, neighborhood_key: str, edges: list[dict[str, Any]], *, warnings: list[str] | None = None ) -> CachedNeighborhood: cache = self.load() neighborhood = cache.neighborhoods.get(neighborhood_key) if neighborhood is None: raise LookupError("Exploration neighborhood not found.") by_id = {str(edge.get("id") or ""): dict(edge) for edge in neighborhood.induced_edges if edge.get("id")} for edge in edges: edge_id = str(edge.get("id") or "") if edge_id: by_id[edge_id] = dict(edge) cache.edges[edge_id] = dict(edge) neighborhood.induced_edges = list(by_id.values()) neighborhood.induced_state = "degraded" if warnings else "complete" neighborhood.warnings = list(dict.fromkeys([*neighborhood.warnings, *(warnings or [])])) neighborhood.updated_at = time.time() neighborhood.expires_at = time.time() + (_ERROR_TTL if warnings else _TOPOLOGY_TTL) cache.neighborhoods[neighborhood_key] = neighborhood self.save(cache) return neighborhood