import json import re import uuid from datetime import datetime from pathlib import Path class SessionManager: def __init__(self, data_dir): self.sessions_dir = Path(data_dir) / "sessions" self.sessions_dir.mkdir(parents=True, exist_ok=True) self._current_id = None self._current_path = None def start_new_session(self, system_message): session_id = uuid.uuid4().hex[:12] now = datetime.now() session = { "id": session_id, "title": "Untitled Session", "created_at": now.isoformat(timespec="seconds"), "updated_at": now.isoformat(timespec="seconds"), "message_count": 0, "messages": [system_message], } file_name = self._build_filename(now, "untitled") file_path = self.sessions_dir / file_name self._current_id = session_id self._current_path = file_path self._write_session(file_path, session) return session def save_messages(self, messages): if self._current_path is None: return session = self._read_session(self._current_path) if session is None: session = { "id": self._current_id or uuid.uuid4().hex[:12], "title": "Restored Session", "created_at": datetime.now().isoformat(timespec="seconds"), } if not self._current_id: self._current_id = session["id"] # Update title from the first user message if still untitled if session["title"] == "Untitled Session": title = self._extract_title(messages) if title: session["title"] = title # Rename the file to include the real title new_path = self._rename_with_title(self._current_path, title) if new_path: self._current_path = new_path session["messages"] = messages session["message_count"] = len([m for m in messages if m["role"] != "system"]) session["updated_at"] = datetime.now().isoformat(timespec="seconds") self._write_session(self._current_path, session) def list_sessions(self): sessions = [] for file_path in sorted(self.sessions_dir.glob("*.json"), reverse=True): session = self._read_session(file_path) if session is None: continue sessions.append({ "path": file_path, "id": session.get("id", ""), "title": session.get("title", "Untitled"), "created_at": session.get("created_at", ""), "updated_at": session.get("updated_at", ""), "message_count": session.get("message_count", 0), }) return sessions def load_session(self, index, system_message): sessions = self.list_sessions() if index < 1 or index > len(sessions): return None entry = sessions[index - 1] session = self._read_session(entry["path"]) if session is None: return None messages = session.get("messages", []) # Ensure the system message is current (in case the prompt was updated) if messages and messages[0].get("role") == "system": messages[0] = system_message else: messages.insert(0, system_message) self._current_id = session.get("id", "") self._current_path = entry["path"] return messages def delete_session(self, index): sessions = self.list_sessions() if index < 1 or index > len(sessions): return None entry = sessions[index - 1] title = entry["title"] try: entry["path"].unlink(missing_ok=True) except OSError: return None # If we deleted the currently active session, clear state if self._current_path == entry["path"]: self._current_id = None self._current_path = None return title def format_session_list(self): sessions = self.list_sessions() if not sessions: return "No saved sessions found." lines = ["", " Saved Sessions", " " + "-" * 50] for index, entry in enumerate(sessions, start=1): count = entry["message_count"] updated = entry["updated_at"][:16].replace("T", " ") title = entry["title"][:50] lines.append(f" {index}. {title}") lines.append(f" {count} messages | Last updated: {updated}") lines.append("") return "\n".join(lines) def _build_filename(self, timestamp, slug): date_part = timestamp.strftime("%Y-%m-%d_%H-%M") safe_slug = self._slugify(slug)[:40] return f"{date_part}_{safe_slug}.json" def _rename_with_title(self, old_path, title): old_path = Path(old_path) if not old_path.exists(): return None # Extract the date prefix from the old filename old_name = old_path.stem parts = old_name.split("_", 3) # e.g. 2026-05-11_18-30_untitled if len(parts) >= 2: date_prefix = f"{parts[0]}_{parts[1]}" else: date_prefix = old_name[:16] safe_title = self._slugify(title)[:40] new_name = f"{date_prefix}_{safe_title}.json" new_path = old_path.parent / new_name if new_path == old_path: return old_path try: old_path.rename(new_path) return new_path except OSError: return old_path def _extract_title(self, messages): for message in messages: if message.get("role") == "user": content = message.get("content", "").strip() if not content: continue # Remove source prefixes like "Source: PDF file (doc.pdf)\n\n..." if content.startswith("Source:"): lines = content.split("\n", 2) title = lines[0].strip() else: title = content # Truncate to something readable if len(title) > 60: title = title[:57] + "..." return title return "" def _slugify(self, text): slug = text.lower().strip() slug = re.sub(r"[^\w\s-]", "", slug) slug = re.sub(r"[\s_]+", "-", slug) slug = slug.strip("-") return slug or "untitled" def _read_session(self, file_path): try: return json.loads(Path(file_path).read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return None def _write_session(self, file_path, session): Path(file_path).write_text( json.dumps(session, ensure_ascii=False, indent=2), encoding="utf-8", )