Update chat_manager.py

chat_manager.py

@@ -1,334 +1,335 @@
-# chat_manager.py - Chat Session Management System
-import json
-import os
-from dataclasses import dataclass, asdict
-from typing import List, Optional, Dict, Any
-from datetime import datetime
-from pathlib import Path
-
-@dataclass
-class ChatMessage:
-    """Individual chat message structure"""
-    message_id: str
-    role: str  # 'user' or 'assistant'
-    content: str
-    timestamp: str
-    rating: Optional[int] = None  # 1 for thumbs up, -1 for thumbs down, None for no rating
-    is_bookmarked: bool = False
-    source_documents: List[str] = None
-    
-    def __post_init__(self):
-        if self.source_documents is None:
-            self.source_documents = []
-
-@dataclass
-class ChatSession:
-    """Chat session data structure"""
-    session_id: str
-    user_id: str
-    title: str
-    created_at: str
-    updated_at: str
-    messages: List[ChatMessage] = None
-    is_archived: bool = False
-    tags: List[str] = None
-    
-    def __post_init__(self):
-        if self.messages is None:
-            self.messages = []
-        if self.tags is None:
-            self.tags = []
-
-class ChatManager:
-    """Manages chat sessions and messages"""
-    
-    def __init__(self, data_dir: str):
-        self.data_dir = Path(data_dir)
-        self.data_dir.mkdir(exist_ok=True)
-        self.sessions_file = self.data_dir / "sessions.json"
-        self.ensure_sessions_file()
-    
-    def ensure_sessions_file(self):
-        """Ensure sessions file exists"""
-        if not self.sessions_file.exists():
-            with open(self.sessions_file, 'w') as f:
-                json.dump({}, f)
-    
-    def create_session(self, user_id: str, title: str = None) -> str:
-        """Create a new chat session"""
-        session_id = str(uuid.uuid4())
-        timestamp = datetime.now().isoformat()
-        
-        if not title:
-            title = f"Chat {datetime.now().strftime('%Y-%m-%d %H:%M')}"
-        
-        session = ChatSession(
-            session_id=session_id,
-            user_id=user_id,
-            title=title,
-            created_at=timestamp,
-            updated_at=timestamp
-        )
-        
-        try:
-            sessions = self.load_all_sessions()
-            sessions[session_id] = asdict(session)
-            
-            with open(self.sessions_file, 'w') as f:
-                json.dump(sessions, f, indent=2)
-            
-            return session_id
-            
-        except Exception as e:
-            raise Exception(f"Failed to create session: {str(e)}")
-    
-    def load_all_sessions(self) -> Dict[str, Dict]:
-        """Load all sessions from storage"""
-        try:
-            with open(self.sessions_file, 'r') as f:
-                return json.load(f)
-        except (FileNotFoundError, json.JSONDecodeError):
-            return {}
-    
-    def get_session(self, session_id: str) -> Optional[ChatSession]:
-        """Get chat session by ID"""
-        sessions = self.load_all_sessions()
-        session_data = sessions.get(session_id)
-        
-        if session_data:
-            # Convert message dictionaries back to ChatMessage objects
-            messages = []
-            for msg_data in session_data.get('messages', []):
-                messages.append(ChatMessage(**msg_data))
-            session_data['messages'] = messages
-            return ChatSession(**session_data)
-        return None
-    
-    def get_user_sessions(self, user_id: str, include_archived: bool = False) -> List[ChatSession]:
-        """Get all sessions for a user"""
-        sessions = self.load_all_sessions()
-        user_sessions = []
-        
-        for session_data in sessions.values():
-            if session_data.get('user_id') == user_id:
-                if include_archived or not session_data.get('is_archived', False):
-                    # Convert message dictionaries back to ChatMessage objects
-                    messages = []
-                    for msg_data in session_data.get('messages', []):
-                        messages.append(ChatMessage(**msg_data))
-                    session_data['messages'] = messages
-                    user_sessions.append(ChatSession(**session_data))
-        
-        # Sort by updated_at descending
-        user_sessions.sort(key=lambda x: x.updated_at, reverse=True)
-        return user_sessions
-    
-    def add_message(self, session_id: str, role: str, content: str, source_documents: List[str] = None) -> str:
-        """Add a message to a chat session"""
-        message_id = str(uuid.uuid4())
-        timestamp = datetime.now().isoformat()
-        
-        message = ChatMessage(
-            message_id=message_id,
-            role=role,
-            content=content,
-            timestamp=timestamp,
-            source_documents=source_documents or []
-        )
-        
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id not in sessions:
-                raise ValueError(f"Session {session_id} not found")
-            
-            # Convert message to dict for storage
-            message_dict = asdict(message)
-            sessions[session_id]['messages'].append(message_dict)
-            sessions[session_id]['updated_at'] = timestamp
-            
-            with open(self.sessions_file, 'w') as f:
-                json.dump(sessions, f, indent=2)
-            
-            return message_id
-            
-        except Exception as e:
-            raise Exception(f"Failed to add message: {str(e)}")
-    
-    def rate_message(self, session_id: str, message_id: str, rating: int) -> bool:
-        """Rate a message (1 for thumbs up, -1 for thumbs down)"""
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id not in sessions:
-                return False
-            
-            for message in sessions[session_id]['messages']:
-                if message['message_id'] == message_id:
-                    message['rating'] = rating
-                    sessions[session_id]['updated_at'] = datetime.now().isoformat()
-                    
-                    with open(self.sessions_file, 'w') as f:
-                        json.dump(sessions, f, indent=2)
-                    
-                    return True
-            
-            return False
-            
-        except Exception:
-            return False
-    
-    def bookmark_message(self, session_id: str, message_id: str, is_bookmarked: bool = True) -> bool:
-        """Bookmark or unbookmark a message"""
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id not in sessions:
-                return False
-            
-            for message in sessions[session_id]['messages']:
-                if message['message_id'] == message_id:
-                    message['is_bookmarked'] = is_bookmarked
-                    sessions[session_id]['updated_at'] = datetime.now().isoformat()
-                    
-                    with open(self.sessions_file, 'w') as f:
-                        json.dump(sessions, f, indent=2)
-                    
-                    return True
-            
-            return False
-            
-        except Exception:
-            return False
-    
-    def get_bookmarked_messages(self, user_id: str) -> List[Dict[str, Any]]:
-        """Get all bookmarked messages for a user"""
-        sessions = self.load_all_sessions()
-        bookmarked = []
-        
-        for session_data in sessions.values():
-            if session_data.get('user_id') == user_id:
-                for message in session_data.get('messages', []):
-                    if message.get('is_bookmarked', False):
-                        bookmarked.append({
-                            'session_id': session_data['session_id'],
-                            'session_title': session_data['title'],
-                            'message': message,
-                            'timestamp': message['timestamp']
-                        })
-        
-        # Sort by timestamp descending
-        bookmarked.sort(key=lambda x: x['timestamp'], reverse=True)
-        return bookmarked
-    
-    def update_session_title(self, session_id: str, title: str) -> bool:
-        """Update session title"""
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id not in sessions:
-                return False
-            
-            sessions[session_id]['title'] = title
-            sessions[session_id]['updated_at'] = datetime.now().isoformat()
-            
-            with open(self.sessions_file, 'w') as f:
-                json.dump(sessions, f, indent=2)
-            
-            return True
-            
-        except Exception:
-            return False
-    
-    def archive_session(self, session_id: str, is_archived: bool = True) -> bool:
-        """Archive or unarchive a session"""
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id not in sessions:
-                return False
-            
-            sessions[session_id]['is_archived'] = is_archived
-            sessions[session_id]['updated_at'] = datetime.now().isoformat()
-            
-            with open(self.sessions_file, 'w') as f:
-                json.dump(sessions, f, indent=2)
-            
-            return True
-            
-        except Exception:
-            return False
-    
-    def delete_session(self, session_id: str) -> bool:
-        """Delete a chat session"""
-        try:
-            sessions = self.load_all_sessions()
-            
-            if session_id in sessions:
-                del sessions[session_id]
-                
-                with open(self.sessions_file, 'w') as f:
-                    json.dump(sessions, f, indent=2)
-                
-                return True
-            return False
-            
-        except Exception:
-            return False
-    
-    def export_chat_history(self, user_id: str, session_id: str = None) -> Dict[str, Any]:
-        """Export chat history for a user or specific session"""
-        if session_id:
-            session = self.get_session(session_id)
-            if session and session.user_id == user_id:
-                return {
-                    'export_type': 'single_session',
-                    'session': asdict(session),
-                    'exported_at': datetime.now().isoformat()
-                }
-        else:
-            sessions = self.get_user_sessions(user_id, include_archived=True)
-            return {
-                'export_type': 'all_sessions',
-                'sessions': [asdict(session) for session in sessions],
-                'exported_at': datetime.now().isoformat(),
-                'total_sessions': len(sessions)
-            }
-        
-        return {}
-    
-    def get_chat_statistics(self, user_id: str) -> Dict[str, Any]:
-        """Get chat statistics for a user"""
-        sessions = self.get_user_sessions(user_id, include_archived=True)
-        
-        total_messages = 0
-        total_user_messages = 0
-        total_assistant_messages = 0
-        bookmarked_count = 0
-        rated_messages = {'positive': 0, 'negative': 0}
-        
-        for session in sessions:
-            total_messages += len(session.messages)
-            for message in session.messages:
-                if message.role == 'user':
-                    total_user_messages += 1
-                else:
-                    total_assistant_messages += 1
-                
-                if message.is_bookmarked:
-                    bookmarked_count += 1
-                
-                if message.rating == 1:
-                    rated_messages['positive'] += 1
-                elif message.rating == -1:
-                    rated_messages['negative'] += 1
-        
-        return {
-            'total_sessions': len(sessions),
-            'total_messages': total_messages,
-            'user_messages': total_user_messages,
-            'assistant_messages': total_assistant_messages,
-            'bookmarked_messages': bookmarked_count,
-            'message_ratings': rated_messages,
-            'average_messages_per_session': total_messages / len(sessions) if sessions else 0
+# chat_manager.py - Chat Session Management System
+import json
+import os
+import uuid
+from dataclasses import dataclass, asdict
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+from pathlib import Path
+
+@dataclass
+class ChatMessage:
+    """Individual chat message structure"""
+    message_id: str
+    role: str  # 'user' or 'assistant'
+    content: str
+    timestamp: str
+    rating: Optional[int] = None  # 1 for thumbs up, -1 for thumbs down, None for no rating
+    is_bookmarked: bool = False
+    source_documents: List[str] = None
+    
+    def __post_init__(self):
+        if self.source_documents is None:
+            self.source_documents = []
+
+@dataclass
+class ChatSession:
+    """Chat session data structure"""
+    session_id: str
+    user_id: str
+    title: str
+    created_at: str
+    updated_at: str
+    messages: List[ChatMessage] = None
+    is_archived: bool = False
+    tags: List[str] = None
+    
+    def __post_init__(self):
+        if self.messages is None:
+            self.messages = []
+        if self.tags is None:
+            self.tags = []
+
+class ChatManager:
+    """Manages chat sessions and messages"""
+    
+    def __init__(self, data_dir: str):
+        self.data_dir = Path(data_dir)
+        self.data_dir.mkdir(exist_ok=True)
+        self.sessions_file = self.data_dir / "sessions.json"
+        self.ensure_sessions_file()
+    
+    def ensure_sessions_file(self):
+        """Ensure sessions file exists"""
+        if not self.sessions_file.exists():
+            with open(self.sessions_file, 'w') as f:
+                json.dump({}, f)
+    
+    def create_session(self, user_id: str, title: str = None) -> str:
+        """Create a new chat session"""
+        session_id = str(uuid.uuid4())
+        timestamp = datetime.now().isoformat()
+        
+        if not title:
+            title = f"Chat {datetime.now().strftime('%Y-%m-%d %H:%M')}"
+        
+        session = ChatSession(
+            session_id=session_id,
+            user_id=user_id,
+            title=title,
+            created_at=timestamp,
+            updated_at=timestamp
+        )
+        
+        try:
+            sessions = self.load_all_sessions()
+            sessions[session_id] = asdict(session)
+            
+            with open(self.sessions_file, 'w') as f:
+                json.dump(sessions, f, indent=2)
+            
+            return session_id
+            
+        except Exception as e:
+            raise Exception(f"Failed to create session: {str(e)}")
+    
+    def load_all_sessions(self) -> Dict[str, Dict]:
+        """Load all sessions from storage"""
+        try:
+            with open(self.sessions_file, 'r') as f:
+                return json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            return {}
+    
+    def get_session(self, session_id: str) -> Optional[ChatSession]:
+        """Get chat session by ID"""
+        sessions = self.load_all_sessions()
+        session_data = sessions.get(session_id)
+        
+        if session_data:
+            # Convert message dictionaries back to ChatMessage objects
+            messages = []
+            for msg_data in session_data.get('messages', []):
+                messages.append(ChatMessage(**msg_data))
+            session_data['messages'] = messages
+            return ChatSession(**session_data)
+        return None
+    
+    def get_user_sessions(self, user_id: str, include_archived: bool = False) -> List[ChatSession]:
+        """Get all sessions for a user"""
+        sessions = self.load_all_sessions()
+        user_sessions = []
+        
+        for session_data in sessions.values():
+            if session_data.get('user_id') == user_id:
+                if include_archived or not session_data.get('is_archived', False):
+                    # Convert message dictionaries back to ChatMessage objects
+                    messages = []
+                    for msg_data in session_data.get('messages', []):
+                        messages.append(ChatMessage(**msg_data))
+                    session_data['messages'] = messages
+                    user_sessions.append(ChatSession(**session_data))
+        
+        # Sort by updated_at descending
+        user_sessions.sort(key=lambda x: x.updated_at, reverse=True)
+        return user_sessions
+    
+    def add_message(self, session_id: str, role: str, content: str, source_documents: List[str] = None) -> str:
+        """Add a message to a chat session"""
+        message_id = str(uuid.uuid4())
+        timestamp = datetime.now().isoformat()
+        
+        message = ChatMessage(
+            message_id=message_id,
+            role=role,
+            content=content,
+            timestamp=timestamp,
+            source_documents=source_documents or []
+        )
+        
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id not in sessions:
+                raise ValueError(f"Session {session_id} not found")
+            
+            # Convert message to dict for storage
+            message_dict = asdict(message)
+            sessions[session_id]['messages'].append(message_dict)
+            sessions[session_id]['updated_at'] = timestamp
+            
+            with open(self.sessions_file, 'w') as f:
+                json.dump(sessions, f, indent=2)
+            
+            return message_id
+            
+        except Exception as e:
+            raise Exception(f"Failed to add message: {str(e)}")
+    
+    def rate_message(self, session_id: str, message_id: str, rating: int) -> bool:
+        """Rate a message (1 for thumbs up, -1 for thumbs down)"""
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id not in sessions:
+                return False
+            
+            for message in sessions[session_id]['messages']:
+                if message['message_id'] == message_id:
+                    message['rating'] = rating
+                    sessions[session_id]['updated_at'] = datetime.now().isoformat()
+                    
+                    with open(self.sessions_file, 'w') as f:
+                        json.dump(sessions, f, indent=2)
+                    
+                    return True
+            
+            return False
+            
+        except Exception:
+            return False
+    
+    def bookmark_message(self, session_id: str, message_id: str, is_bookmarked: bool = True) -> bool:
+        """Bookmark or unbookmark a message"""
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id not in sessions:
+                return False
+            
+            for message in sessions[session_id]['messages']:
+                if message['message_id'] == message_id:
+                    message['is_bookmarked'] = is_bookmarked
+                    sessions[session_id]['updated_at'] = datetime.now().isoformat()
+                    
+                    with open(self.sessions_file, 'w') as f:
+                        json.dump(sessions, f, indent=2)
+                    
+                    return True
+            
+            return False
+            
+        except Exception:
+            return False
+    
+    def get_bookmarked_messages(self, user_id: str) -> List[Dict[str, Any]]:
+        """Get all bookmarked messages for a user"""
+        sessions = self.load_all_sessions()
+        bookmarked = []
+        
+        for session_data in sessions.values():
+            if session_data.get('user_id') == user_id:
+                for message in session_data.get('messages', []):
+                    if message.get('is_bookmarked', False):
+                        bookmarked.append({
+                            'session_id': session_data['session_id'],
+                            'session_title': session_data['title'],
+                            'message': message,
+                            'timestamp': message['timestamp']
+                        })
+        
+        # Sort by timestamp descending
+        bookmarked.sort(key=lambda x: x['timestamp'], reverse=True)
+        return bookmarked
+    
+    def update_session_title(self, session_id: str, title: str) -> bool:
+        """Update session title"""
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id not in sessions:
+                return False
+            
+            sessions[session_id]['title'] = title
+            sessions[session_id]['updated_at'] = datetime.now().isoformat()
+            
+            with open(self.sessions_file, 'w') as f:
+                json.dump(sessions, f, indent=2)
+            
+            return True
+            
+        except Exception:
+            return False
+    
+    def archive_session(self, session_id: str, is_archived: bool = True) -> bool:
+        """Archive or unarchive a session"""
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id not in sessions:
+                return False
+            
+            sessions[session_id]['is_archived'] = is_archived
+            sessions[session_id]['updated_at'] = datetime.now().isoformat()
+            
+            with open(self.sessions_file, 'w') as f:
+                json.dump(sessions, f, indent=2)
+            
+            return True
+            
+        except Exception:
+            return False
+    
+    def delete_session(self, session_id: str) -> bool:
+        """Delete a chat session"""
+        try:
+            sessions = self.load_all_sessions()
+            
+            if session_id in sessions:
+                del sessions[session_id]
+                
+                with open(self.sessions_file, 'w') as f:
+                    json.dump(sessions, f, indent=2)
+                
+                return True
+            return False
+            
+        except Exception:
+            return False
+    
+    def export_chat_history(self, user_id: str, session_id: str = None) -> Dict[str, Any]:
+        """Export chat history for a user or specific session"""
+        if session_id:
+            session = self.get_session(session_id)
+            if session and session.user_id == user_id:
+                return {
+                    'export_type': 'single_session',
+                    'session': asdict(session),
+                    'exported_at': datetime.now().isoformat()
+                }
+        else:
+            sessions = self.get_user_sessions(user_id, include_archived=True)
+            return {
+                'export_type': 'all_sessions',
+                'sessions': [asdict(session) for session in sessions],
+                'exported_at': datetime.now().isoformat(),
+                'total_sessions': len(sessions)
+            }
+        
+        return {}
+    
+    def get_chat_statistics(self, user_id: str) -> Dict[str, Any]:
+        """Get chat statistics for a user"""
+        sessions = self.get_user_sessions(user_id, include_archived=True)
+        
+        total_messages = 0
+        total_user_messages = 0
+        total_assistant_messages = 0
+        bookmarked_count = 0
+        rated_messages = {'positive': 0, 'negative': 0}
+        
+        for session in sessions:
+            total_messages += len(session.messages)
+            for message in session.messages:
+                if message.role == 'user':
+                    total_user_messages += 1
+                else:
+                    total_assistant_messages += 1
+                
+                if message.is_bookmarked:
+                    bookmarked_count += 1
+                
+                if message.rating == 1:
+                    rated_messages['positive'] += 1
+                elif message.rating == -1:
+                    rated_messages['negative'] += 1
+        
+        return {
+            'total_sessions': len(sessions),
+            'total_messages': total_messages,
+            'user_messages': total_user_messages,
+            'assistant_messages': total_assistant_messages,
+            'bookmarked_messages': bookmarked_count,
+            'message_ratings': rated_messages,
+            'average_messages_per_session': total_messages / len(sessions) if sessions else 0
         }