| """ |
| episodic_store.py |
| |
| Lazy loader for all_sessions.json. Keeps the full raw-turn data in memory |
| and returns sessions on demand by session ID, avoiding loading all sessions |
| into ChatHistory upfront. |
| """ |
|
|
| import json |
| from datetime import datetime |
| from typing import Dict, List, Optional |
|
|
|
|
| class EpisodicMemoryStore: |
| def __init__(self, all_sessions_path: str): |
| print(f"[EpisodicMemoryStore] Loading {all_sessions_path} ...") |
| with open(all_sessions_path) as f: |
| self._data: Dict[str, List] = json.load(f) |
| print(f"[EpisodicMemoryStore] Loaded {len(self._data)} sessions.") |
|
|
| @staticmethod |
| def _parse_date(date_str: str) -> datetime: |
| """Convert '2023/04/10 (Mon) 17:50' to datetime.""" |
| date_part = date_str.split('(')[0].strip() |
| time_part = date_str.split(')')[-1].strip() |
| return datetime.strptime(date_part + time_part, "%Y/%m/%d%H:%M") |
|
|
| def get_raw_sessions( |
| self, |
| sess_ids: List[str], |
| date_lookup: Dict[str, str], |
| topic_lookup: Optional[Dict[str, List[str]]] = None, |
| ) -> List[dict]: |
| """ |
| Return a list of session dicts compatible with ChatHistory(sessions=...). |
| |
| Args: |
| sess_ids: Ordered list of session IDs to load. |
| date_lookup: Mapping sess_id -> date string (e.g. '2023/04/10 (Mon) 17:50'). |
| topic_lookup: Optional mapping sess_id -> list of topic strings. |
| |
| Returns: |
| List of session dicts, each with keys: |
| session_id, session_date, session, topic, timestamp |
| """ |
| sessions = [] |
| for sid in sess_ids: |
| if sid not in self._data: |
| continue |
| date_str = date_lookup.get(sid, "") |
| try: |
| ts = self._parse_date(date_str) if date_str else datetime.min |
| except Exception: |
| ts = datetime.min |
| sessions.append({ |
| "session_id": sid, |
| "session_date": date_str, |
| "session": self._data[sid], |
| "topic": topic_lookup.get(sid, []) if topic_lookup else [], |
| "timestamp": ts, |
| }) |
| return sessions |
|
|
