|
import json
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
from pr.core.logging import get_logger
|
|
|
|
logger = get_logger('session')
|
|
|
|
SESSIONS_DIR = os.path.expanduser("~/.assistant_sessions")
|
|
|
|
|
|
class SessionManager:
|
|
|
|
def __init__(self):
|
|
os.makedirs(SESSIONS_DIR, exist_ok=True)
|
|
|
|
def save_session(self, name: str, messages: List[Dict], metadata: Optional[Dict] = None) -> bool:
|
|
try:
|
|
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
|
|
|
|
session_data = {
|
|
'name': name,
|
|
'created_at': datetime.now().isoformat(),
|
|
'messages': messages,
|
|
'metadata': metadata or {}
|
|
}
|
|
|
|
with open(session_file, 'w') as f:
|
|
json.dump(session_data, f, indent=2)
|
|
|
|
logger.info(f"Session saved: {name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error saving session {name}: {e}")
|
|
return False
|
|
|
|
def load_session(self, name: str) -> Optional[Dict]:
|
|
try:
|
|
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
|
|
|
|
if not os.path.exists(session_file):
|
|
logger.warning(f"Session not found: {name}")
|
|
return None
|
|
|
|
with open(session_file, 'r') as f:
|
|
session_data = json.load(f)
|
|
|
|
logger.info(f"Session loaded: {name}")
|
|
return session_data
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error loading session {name}: {e}")
|
|
return None
|
|
|
|
def list_sessions(self) -> List[Dict]:
|
|
sessions = []
|
|
|
|
try:
|
|
for filename in os.listdir(SESSIONS_DIR):
|
|
if filename.endswith('.json'):
|
|
filepath = os.path.join(SESSIONS_DIR, filename)
|
|
try:
|
|
with open(filepath, 'r') as f:
|
|
data = json.load(f)
|
|
|
|
sessions.append({
|
|
'name': data.get('name', filename[:-5]),
|
|
'created_at': data.get('created_at', 'unknown'),
|
|
'message_count': len(data.get('messages', [])),
|
|
'metadata': data.get('metadata', {})
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"Error reading session file {filename}: {e}")
|
|
|
|
sessions.sort(key=lambda x: x['created_at'], reverse=True)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error listing sessions: {e}")
|
|
|
|
return sessions
|
|
|
|
def delete_session(self, name: str) -> bool:
|
|
try:
|
|
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
|
|
|
|
if not os.path.exists(session_file):
|
|
logger.warning(f"Session not found: {name}")
|
|
return False
|
|
|
|
os.remove(session_file)
|
|
logger.info(f"Session deleted: {name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error deleting session {name}: {e}")
|
|
return False
|
|
|
|
def export_session(self, name: str, output_path: str, format: str = 'json') -> bool:
|
|
session_data = self.load_session(name)
|
|
if not session_data:
|
|
return False
|
|
|
|
try:
|
|
if format == 'json':
|
|
with open(output_path, 'w') as f:
|
|
json.dump(session_data, f, indent=2)
|
|
|
|
elif format == 'markdown':
|
|
with open(output_path, 'w') as f:
|
|
f.write(f"# Session: {name}\n\n")
|
|
f.write(f"Created: {session_data['created_at']}\n\n")
|
|
f.write("---\n\n")
|
|
|
|
for msg in session_data['messages']:
|
|
role = msg.get('role', 'unknown')
|
|
content = msg.get('content', '')
|
|
|
|
f.write(f"## {role.capitalize()}\n\n")
|
|
f.write(f"{content}\n\n")
|
|
f.write("---\n\n")
|
|
|
|
elif format == 'txt':
|
|
with open(output_path, 'w') as f:
|
|
f.write(f"Session: {name}\n")
|
|
f.write(f"Created: {session_data['created_at']}\n")
|
|
f.write("=" * 80 + "\n\n")
|
|
|
|
for msg in session_data['messages']:
|
|
role = msg.get('role', 'unknown')
|
|
content = msg.get('content', '')
|
|
|
|
f.write(f"[{role.upper()}]\n")
|
|
f.write(f"{content}\n")
|
|
f.write("-" * 80 + "\n\n")
|
|
|
|
else:
|
|
logger.error(f"Unsupported export format: {format}")
|
|
return False
|
|
|
|
logger.info(f"Session exported to {output_path}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error exporting session: {e}")
|
|
return False
|