Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| from huggingface_hub import HfApi, hf_hub_download | |
| from src.schemas import DatasetRecord, Opportunity, ResearchResult, record_from_opportunity | |
| class DatasetStoreError(RuntimeError): | |
| pass | |
| class DatasetStore: | |
| def __init__( | |
| self, | |
| token: str | None = None, | |
| repo_id: str | None = None, | |
| filename: str = "scholarscope_sessions.jsonl", | |
| ) -> None: | |
| self.token = os.getenv("HF_TOKEN", "") if token is None else token | |
| self.repo_id = os.getenv("HF_DATASET_REPO", "") if repo_id is None else repo_id | |
| self.filename = filename | |
| self.api = HfApi(token=self.token or None) | |
| self._memory_records: list[DatasetRecord] = [] | |
| def enabled(self) -> bool: | |
| return bool(self.token and self.repo_id) | |
| def record_search(self, result: ResearchResult) -> None: | |
| if not result.opportunities: | |
| return | |
| top_opportunity = max(result.opportunities, key=lambda opportunity: opportunity.match_score) | |
| self.append_records([record_from_opportunity(result, top_opportunity)]) | |
| def save_opportunity(self, result: ResearchResult, opportunity: Opportunity) -> DatasetRecord: | |
| record = record_from_opportunity(result, opportunity) | |
| self.append_records([record]) | |
| return record | |
| def append_records(self, records: list[DatasetRecord]) -> None: | |
| if not records: | |
| return | |
| existing = self._download_existing() if self.enabled else [] | |
| existing_records = self._memory_records if not self.enabled else self._parse_records(existing) | |
| seen = {self._record_key(record) for record in existing_records} | |
| relevant: list[DatasetRecord] = [] | |
| for record in records: | |
| key = self._record_key(record) | |
| if key in seen: | |
| continue | |
| seen.add(key) | |
| relevant.append(record) | |
| if not relevant: | |
| return | |
| if not self.enabled: | |
| self._memory_records.extend(relevant) | |
| return | |
| lines = existing + [record.model_dump_json() for record in relevant] | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| local_path = Path(temp_dir) / self.filename | |
| local_path.write_text("\n".join(lines) + "\n", encoding="utf-8") | |
| try: | |
| self.api.create_repo( | |
| repo_id=self.repo_id, | |
| repo_type="dataset", | |
| exist_ok=True, | |
| private=True, | |
| ) | |
| self.api.upload_file( | |
| path_or_fileobj=str(local_path), | |
| path_in_repo=self.filename, | |
| repo_id=self.repo_id, | |
| repo_type="dataset", | |
| commit_message="Update ScholarScope research records", | |
| ) | |
| except Exception as exc: # noqa: BLE001 - convert SDK errors into UI-safe message | |
| raise DatasetStoreError(f"Could not write to Hugging Face Dataset repo: {exc}") from exc | |
| def recent_records(self, limit: int = 20) -> list[DatasetRecord]: | |
| lines = self._download_existing() if self.enabled else [r.model_dump_json() for r in self._memory_records] | |
| records: list[DatasetRecord] = [] | |
| for line in lines[-limit:]: | |
| try: | |
| records.append(DatasetRecord(**json.loads(line))) | |
| except (json.JSONDecodeError, ValueError): | |
| continue | |
| return records | |
| def search_records(self, query: str, limit: int = 50) -> list[DatasetRecord]: | |
| """Return records whose serialized content contains the query (case-insensitive).""" | |
| lines = self._download_existing() if self.enabled else [r.model_dump_json() for r in self._memory_records] | |
| needle = (query or "").strip().lower() | |
| records: list[DatasetRecord] = [] | |
| for line in lines: | |
| if needle and needle not in line.lower(): | |
| continue | |
| try: | |
| records.append(DatasetRecord(**json.loads(line))) | |
| except (json.JSONDecodeError, ValueError): | |
| continue | |
| return records[-limit:] | |
| def _record_key(record: DatasetRecord) -> tuple[str, str]: | |
| return (record.session_id, record.source_url.rstrip("/").lower()) | |
| def _parse_records(lines: list[str]) -> list[DatasetRecord]: | |
| records: list[DatasetRecord] = [] | |
| for line in lines: | |
| try: | |
| records.append(DatasetRecord(**json.loads(line))) | |
| except (json.JSONDecodeError, ValueError): | |
| continue | |
| return records | |
| def _download_existing(self) -> list[str]: | |
| if not self.enabled: | |
| return [] | |
| try: | |
| path = hf_hub_download( | |
| repo_id=self.repo_id, | |
| filename=self.filename, | |
| repo_type="dataset", | |
| token=self.token, | |
| ) | |
| except Exception: | |
| return [] | |
| return [line for line in Path(path).read_text(encoding="utf-8").splitlines() if line.strip()] | |