Spaces:
Paused
Paused
File size: 10,028 Bytes
d5f2660 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 |
"""File-based session management with automatic cleanup"""
import os
import json
import shutil
import time
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any, List
from .logger import setup_logger
from .config import HISTORY_EXCHANGES, SESSION_TTL
logger = setup_logger(__name__)
# Configuration
HISTORY_DIR = Path("History")
class SessionManager:
"""Manages session history with file-based storage"""
def __init__(self):
"""Initialize session manager and create History directory"""
HISTORY_DIR.mkdir(exist_ok=True)
logger.info(f"Session manager initialized | History: {HISTORY_DIR}")
def _get_session_dir(self, session_id: str) -> Path:
"""Get session directory path"""
return HISTORY_DIR / f"Session_{session_id}"
def _get_session_json(self, session_id: str) -> Path:
"""Get session JSON file path"""
return self._get_session_dir(session_id) / "history.json"
def _get_images_dir(self, session_id: str) -> Path:
"""Get session images directory path"""
return self._get_session_dir(session_id) / "images"
def create_session(self, session_id: str) -> None:
"""Create new session directory structure"""
session_dir = self._get_session_dir(session_id)
images_dir = self._get_images_dir(session_id)
session_dir.mkdir(exist_ok=True)
images_dir.mkdir(exist_ok=True)
# Initialize history JSON
now = datetime.now().isoformat()
history_data = {
"session_id": session_id,
"created_at": now,
"last_activity": now,
"messages": []
}
with open(self._get_session_json(session_id), 'w') as f:
json.dump(history_data, f, indent=2)
logger.info(f"🟢 NEW SESSION: {session_id[:8]}")
def session_exists(self, session_id: str) -> bool:
"""Check if session exists"""
return self._get_session_dir(session_id).exists()
def get_session_history(self, session_id: str) -> List[Dict[str, Any]]:
"""Load session history from JSON"""
if not self.session_exists(session_id):
self.create_session(session_id)
return []
json_path = self._get_session_json(session_id)
with open(json_path, 'r') as f:
data = json.load(f)
return data.get("messages", [])
def add_message(self, session_id: str, role: str, content: str,
image_path: Optional[str] = None) -> None:
"""
Add message to session history with sliding window (3 exchanges)
Args:
session_id: Session identifier
role: 'user' or 'assistant'
content: Message text
image_path: Optional path to image file
"""
if not self.session_exists(session_id):
self.create_session(session_id)
# Load current history
json_path = self._get_session_json(session_id)
with open(json_path, 'r') as f:
data = json.load(f)
# Prepare and add message
message = self._create_message(session_id, role, content, image_path, len(data['messages']))
data["messages"].append(message)
data["last_activity"] = datetime.now().isoformat()
# Trim to keep only last 3 exchanges (6 messages)
self._trim_history(session_id, data)
# Save updated history
with open(json_path, 'w') as f:
json.dump(data, f, indent=2)
def _create_message(self, session_id: str, role: str, content: str,
image_path: Optional[str], msg_count: int) -> Dict[str, Any]:
"""Create a message dict with optional image"""
message = {
"role": role,
"content": content,
"timestamp": datetime.now().isoformat(),
"image": None
}
# Handle image if provided
if image_path and os.path.exists(image_path):
image_filename = f"msg_{msg_count}_{int(time.time())}.png"
dest_path = self._get_images_dir(session_id) / image_filename
shutil.copy2(image_path, dest_path)
message["image"] = image_filename
return message
def _trim_history(self, session_id: str, data: Dict[str, Any]) -> None:
"""Trim message history to max exchanges and delete old images"""
max_messages = HISTORY_EXCHANGES * 2
if len(data["messages"]) <= max_messages:
return
# Calculate how many to remove
num_to_remove = len(data["messages"]) - max_messages
removed = data["messages"][:num_to_remove]
data["messages"] = data["messages"][num_to_remove:]
# Delete old images
for msg in removed:
if msg.get("image"):
old_img = self._get_images_dir(session_id) / msg["image"]
if old_img.exists():
old_img.unlink()
logger.info(f"🗑️ Image dropped: {msg['image']}")
def get_messages_for_context(self, session_id: str) -> List[Dict[str, Any]]:
"""Get messages for conversation context with full image paths"""
messages = self.get_session_history(session_id)
recent = messages # Already trimmed to 3 exchanges in add_message
# Convert image filenames to full paths
for msg in recent:
if msg.get("image"):
msg["image_path"] = str(self._get_images_dir(session_id) / msg["image"])
return recent
def update_activity(self, session_id: str) -> None:
"""Update session last activity timestamp"""
if not self.session_exists(session_id):
return
json_path = self._get_session_json(session_id)
with open(json_path, 'r') as f:
data = json.load(f)
data["last_activity"] = datetime.now().isoformat()
with open(json_path, 'w') as f:
json.dump(data, f, indent=2)
def clear_session(self, session_id: str) -> None:
"""Clear session history (keep structure)"""
if not self.session_exists(session_id):
return
# Clear images
images_dir = self._get_images_dir(session_id)
for img in images_dir.glob("*"):
img.unlink()
# Reset history
json_path = self._get_session_json(session_id)
with open(json_path, 'r') as f:
data = json.load(f)
data["messages"] = []
data["last_activity"] = datetime.now().isoformat()
with open(json_path, 'w') as f:
json.dump(data, f, indent=2)
def delete_session(self, session_id: str) -> None:
"""Delete entire session directory"""
session_dir = self._get_session_dir(session_id)
if session_dir.exists():
shutil.rmtree(session_dir)
logger.info(f"🔴 EXPIRED: {session_id[:8]}")
def cleanup_expired_sessions(self) -> int:
"""
Delete sessions that exceeded TTL (5 minutes)
Returns number of sessions deleted
"""
if not HISTORY_DIR.exists():
return 0
current_time = datetime.now()
deleted_count = 0
for session_dir in HISTORY_DIR.glob("Session_*"):
if self._should_delete_session(session_dir, current_time):
deleted_count += 1
return deleted_count
def _should_delete_session(self, session_dir: Path, current_time: datetime) -> bool:
"""Check if session should be deleted and delete if expired"""
json_path = session_dir / "history.json"
if not json_path.exists():
return False
try:
with open(json_path, 'r') as f:
data = json.load(f)
last_activity = datetime.fromisoformat(data["last_activity"])
idle_seconds = (current_time - last_activity).total_seconds()
if idle_seconds > SESSION_TTL:
session_id = data["session_id"]
self.delete_session(session_id)
return True
except Exception as e:
logger.error(f"Error checking session {session_dir.name}: {e}")
return False
def get_active_sessions(self) -> List[Dict[str, Any]]:
"""Get list of active sessions (within TTL)"""
if not HISTORY_DIR.exists():
return []
current_time = datetime.now()
active = []
for session_dir in HISTORY_DIR.glob("Session_*"):
json_path = session_dir / "history.json"
if not json_path.exists():
continue
try:
with open(json_path, 'r') as f:
data = json.load(f)
last_activity = datetime.fromisoformat(data["last_activity"])
idle_seconds = (current_time - last_activity).total_seconds()
if idle_seconds <= SESSION_TTL:
active.append({
"session_id": data["session_id"][:8],
"created": datetime.fromisoformat(data["created_at"]).strftime("%H:%M:%S"),
"last_active": last_activity.strftime("%H:%M:%S"),
"messages": len(data["messages"]),
"idle_seconds": int(idle_seconds)
})
except Exception as e:
logger.error(f"Error reading session {session_dir.name}: {e}")
return active
# Global instance
session_manager = SessionManager()
|