Update services/session_manager.py

services/session_manager.py

@@ -1,412 +1,412 @@
-# services/session_manager.py
-"""
-MongoDB-based user session management with pipeline tracking
-"""
-import os
-import uuid
-from datetime import datetime, timedelta
-from typing import Optional, Dict, Any, List
-from pymongo import MongoClient
-from pymongo.errors import DuplicateKeyError, ConnectionFailure
-
-
-class SessionManager:
-    """
-    Manages user sessions in MongoDB with pipeline execution tracking
-    """
-    
-    def __init__(self):
-        """Initialize MongoDB connection for sessions"""
-        self.mongo_uri = os.getenv("MONGODB_URI")
-        self.db_name = os.getenv("MONGODB_DB", "point9")
-        self.collection_name = "user-sessions"  # New collection for sessions
-        self.pipelines_collection_name = "pipeline-executions"  # Track pipeline runs
-        
-        self.client = None
-        self.db = None
-        self.sessions_col = None
-        self.pipelines_col = None
-        
-        self._connect()
-    
-    def _connect(self):
-        """Establish MongoDB connection"""
-        if not self.mongo_uri:
-            print("⚠️  MongoDB URI not configured - session persistence disabled")
-            return
-        
-        try:
-            self.client = MongoClient(self.mongo_uri, serverSelectionTimeoutMS=5000)
-            self.client.admin.command("ping")  # Test connection
-            
-            self.db = self.client[self.db_name]
-            self.sessions_col = self.db[self.collection_name]
-            self.pipelines_col = self.db[self.pipelines_collection_name]
-            
-            # Create indexes
-            self.sessions_col.create_index("session_id", unique=True)
-            self.sessions_col.create_index("created_at")
-            self.sessions_col.create_index("last_activity")
-            
-            self.pipelines_col.create_index("session_id")
-            self.pipelines_col.create_index("executed_at")
-            self.pipelines_col.create_index("pipeline_name")
-            
-            print(f"✅ MongoDB session manager connected: {self.db_name}.{self.collection_name}")
-            
-        except ConnectionFailure as e:
-            print(f"❌ MongoDB connection failed: {e}")
-            self.client = None
-    
-    def create_session(
-        self,
-        user_id: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> str:
-        """
-        Create a new user session
-        
-        Args:
-            user_id: Optional user identifier
-            metadata: Additional session metadata
-        
-        Returns:
-            session_id: Unique session identifier
-        """
-        session_id = str(uuid.uuid4())
-        
-        session_data = {
-            "session_id": session_id,
-            "user_id": user_id,
-            "created_at": datetime.now(),
-            "last_activity": datetime.now(),
-            "current_file": None,
-            "proposed_pipeline": None,
-            "state": "initial",  # initial, pipeline_proposed, executing, completed
-            "conversation_history": [],
-            "pipeline_executions": [],
-            "metadata": metadata or {},
-            "stats": {
-                "total_messages": 0,
-                "total_pipelines_executed": 0,
-                "total_tokens_used": 0
-            }
-        }
-        
-        if self.sessions_col is not None:
-            try:
-                self.sessions_col.insert_one(session_data)
-                print(f"✅ Session created in MongoDB: {session_id}")
-            except Exception as e:
-                print(f"⚠️  Failed to save session to MongoDB: {e}")
-        
-        return session_id
-    
-    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
-        """
-        Retrieve session by ID
-        
-        Args:
-            session_id: Session identifier
-        
-        Returns:
-            Session data or None if not found
-        """
-        if self.sessions_col is None:
-            return None
-        
-        try:
-            session = self.sessions_col.find_one({"session_id": session_id})
-            if session:
-                # Update last activity
-                self.sessions_col.update_one(
-                    {"session_id": session_id},
-                    {"$set": {"last_activity": datetime.now()}}
-                )
-                # Remove MongoDB _id field
-                session.pop("_id", None)
-                return session
-            return None
-        except Exception as e:
-            print(f"⚠️  Error retrieving session: {e}")
-            return None
-    
-    def update_session(
-        self,
-        session_id: str,
-        updates: Dict[str, Any]
-    ) -> bool:
-        """
-        Update session data
-        
-        Args:
-            session_id: Session identifier
-            updates: Dictionary of fields to update
-        
-        Returns:
-            True if successful, False otherwise
-        """
-        if self.sessions_col is None:
-            return False
-        
-        try:
-            updates["last_activity"] = datetime.now()
-            
-            result = self.sessions_col.update_one(
-                {"session_id": session_id},
-                {"$set": updates}
-            )
-            
-            return result.modified_count > 0
-        except Exception as e:
-            print(f"⚠️  Error updating session: {e}")
-            return False
-    
-    def add_message(
-        self,
-        session_id: str,
-        role: str,
-        content: str,
-        metadata: Optional[Dict[str, Any]] = None
-    ) -> bool:
-        """
-        Add a message to conversation history
-        
-        Args:
-            session_id: Session identifier
-            role: Message role (user, assistant, system)
-            content: Message content
-            metadata: Additional message metadata
-        
-        Returns:
-            True if successful
-        """
-        if self.sessions_col is None:
-            return False
-        
-        try:
-            message = {
-                "role": role,
-                "content": content,
-                "timestamp": datetime.now(),
-                "metadata": metadata or {}
-            }
-            
-            self.sessions_col.update_one(
-                {"session_id": session_id},
-                {
-                    "$push": {"conversation_history": message},
-                    "$inc": {"stats.total_messages": 1},
-                    "$set": {"last_activity": datetime.now()}
-                }
-            )
-            
-            return True
-        except Exception as e:
-            print(f"⚠️  Error adding message: {e}")
-            return False
-    
-    def save_pipeline_execution(
-        self,
-        session_id: str,
-        pipeline: Dict[str, Any],
-        result: Dict[str, Any],
-        file_path: Optional[str] = None,
-        executor: str = "unknown"
-    ) -> bool:
-        """
-        Save pipeline execution to dedicated collection
-        
-        Args:
-            session_id: Session identifier
-            pipeline: Pipeline configuration
-            result: Execution result
-            file_path: File that was processed
-            executor: Which executor was used (bedrock, crewai, gemini)
-        
-        Returns:
-            True if successful
-        """
-        if self.pipelines_col is None:
-            return False
-        
-        try:
-            execution_data = {
-                "execution_id": str(uuid.uuid4()),
-                "session_id": session_id,
-                "pipeline_name": pipeline.get("pipeline_name"),
-                "pipeline_config": pipeline,
-                "result": result,
-                "file_path": file_path,
-                "executor": executor,
-                "executed_at": datetime.now(),
-                "duration_seconds": result.get("summary", {}).get("total_duration_seconds"),
-                "status": result.get("status", "unknown"),
-                "components_executed": len(pipeline.get("components", []))
-            }
-            
-            self.pipelines_col.insert_one(execution_data)
-            
-            # Update session stats
-            self.sessions_col.update_one(
-                {"session_id": session_id},
-                {
-                    "$inc": {"stats.total_pipelines_executed": 1},
-                    "$push": {"pipeline_executions": execution_data["execution_id"]}
-                }
-            )
-            
-            print(f"✅ Pipeline execution saved: {execution_data['execution_id']}")
-            return True
-            
-        except Exception as e:
-            print(f"⚠️  Error saving pipeline execution: {e}")
-            return False
-    
-    def get_session_history(
-        self,
-        session_id: str,
-        limit: int = 50
-    ) -> List[Dict[str, Any]]:
-        """
-        Get conversation history for a session
-        
-        Args:
-            session_id: Session identifier
-            limit: Maximum number of messages to return
-        
-        Returns:
-            List of messages
-        """
-        session = self.get_session(session_id)
-        if not session:
-            return []
-        
-        history = session.get("conversation_history", [])
-        return history[-limit:] if len(history) > limit else history
-    
-    def get_pipeline_executions(
-        self,
-        session_id: Optional[str] = None,
-        limit: int = 10
-    ) -> List[Dict[str, Any]]:
-        """
-        Get pipeline execution history
-        
-        Args:
-            session_id: Optional session filter
-            limit: Maximum number of executions to return
-        
-        Returns:
-            List of pipeline executions
-        """
-        if self.pipelines_col is None:
-            return []
-        
-        try:
-            query = {"session_id": session_id} if session_id else {}
-            
-            executions = self.pipelines_col.find(query).sort("executed_at", -1).limit(limit)
-            
-            result = []
-            for exec_doc in executions:
-                exec_doc.pop("_id", None)
-                # Convert datetime to ISO string
-                if "executed_at" in exec_doc and isinstance(exec_doc["executed_at"], datetime):
-                    exec_doc["executed_at"] = exec_doc["executed_at"].isoformat()
-                result.append(exec_doc)
-            
-            return result
-            
-        except Exception as e:
-            print(f"⚠️  Error retrieving pipeline executions: {e}")
-            return []
-    
-    def cleanup_old_sessions(self, max_age_hours: int = 24) -> int:
-        """
-        Remove sessions older than max_age_hours
-        
-        Args:
-            max_age_hours: Maximum session age in hours
-        
-        Returns:
-            Number of sessions removed
-        """
-        if self.sessions_col is None:
-            return 0
-        
-        try:
-            cutoff = datetime.now() - timedelta(hours=max_age_hours)
-            
-            result = self.sessions_col.delete_many({
-                "last_activity": {"$lt": cutoff}
-            })
-            
-            count = result.deleted_count
-            if count > 0:
-                print(f"🧹 Cleaned up {count} old sessions")
-            
-            return count
-            
-        except Exception as e:
-            print(f"⚠️  Error cleaning up sessions: {e}")
-            return 0
-    
-    def get_session_stats(self, session_id: str) -> Dict[str, Any]:
-        """
-        Get statistics for a session
-        
-        Args:
-            session_id: Session identifier
-        
-        Returns:
-            Session statistics
-        """
-        session = self.get_session(session_id)
-        if not session:
-            return {}
-        
-        return {
-            "session_id": session_id,
-            "created_at": session.get("created_at"),
-            "last_activity": session.get("last_activity"),
-            "total_messages": session.get("stats", {}).get("total_messages", 0),
-            "total_pipelines_executed": session.get("stats", {}).get("total_pipelines_executed", 0),
-            "conversation_length": len(session.get("conversation_history", [])),
-            "state": session.get("state", "unknown")
-        }
-    
-    def close(self):
-        """Close MongoDB connection"""
-        if self.client:
-            self.client.close()
-            print("🔒 MongoDB connection closed")
-
-
-# Global session manager instance
-session_manager = SessionManager()
-
-
-if __name__ == "__main__":
-    # Test session manager
-    print("Testing Session Manager...")
-    
-    # Create session
-    sid = session_manager.create_session(user_id="test_user")
-    print(f"Created session: {sid}")
-    
-    # Add messages
-    session_manager.add_message(sid, "user", "Hello!")
-    session_manager.add_message(sid, "assistant", "Hi! How can I help?")
-    
-    # Get session
-    session = session_manager.get_session(sid)
-    print(f"Session data: {session}")
-    
-    # Get history
-    history = session_manager.get_session_history(sid)
-    print(f"History: {history}")
-    
-    # Get stats
-    stats = session_manager.get_session_stats(sid)
-    print(f"Stats: {stats}")
+# services/session_manager.py
+"""
+MongoDB-based user session management with pipeline tracking
+"""
+import os
+import uuid
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Any, List
+from pymongo import MongoClient
+from pymongo.errors import DuplicateKeyError, ConnectionFailure
+
+
+class SessionManager:
+    """
+    Manages user sessions in MongoDB with pipeline execution tracking
+    """
+    
+    def __init__(self):
+        """Initialize MongoDB connection for sessions"""
+        self.mongo_uri = os.getenv("MONGODB_URI")
+        self.db_name = os.getenv("MONGODB_DB", "point9")
+        self.collection_name = "user-sessions"  # New collection for sessions
+        self.pipelines_collection_name = "pipeline-executions"  # Track pipeline runs
+        
+        self.client = None
+        self.db = None
+        self.sessions_col = None
+        self.pipelines_col = None
+        
+        self._connect()
+    
+    def _connect(self):
+        """Establish MongoDB connection"""
+        if not self.mongo_uri:
+            print("⚠️  MongoDB URI not configured - session persistence disabled")
+            return
+        
+        try:
+            self.client = MongoClient(self.mongo_uri, serverSelectionTimeoutMS=5000)
+            self.client.admin.command("ping")  # Test connection
+            
+            self.db = self.client[self.db_name]
+            self.sessions_col = self.db[self.collection_name]
+            self.pipelines_col = self.db[self.pipelines_collection_name]
+            
+            # Create indexes
+            self.sessions_col.create_index("session_id", unique=True)
+            self.sessions_col.create_index("created_at")
+            self.sessions_col.create_index("last_activity")
+            
+            self.pipelines_col.create_index("session_id")
+            self.pipelines_col.create_index("executed_at")
+            self.pipelines_col.create_index("pipeline_name")
+            
+            print(f"✅ MongoDB session manager connected: {self.db_name}.{self.collection_name}")
+            
+        except ConnectionFailure as e:
+            print(f"❌ MongoDB connection failed: {e}")
+            self.client = None
+    
+    def create_session(
+        self,
+        user_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> str:
+        """
+        Create a new user session
+        
+        Args:
+            user_id: Optional user identifier
+            metadata: Additional session metadata
+        
+        Returns:
+            session_id: Unique session identifier
+        """
+        session_id = str(uuid.uuid4())
+        
+        session_data = {
+            "session_id": session_id,
+            "user_id": user_id,
+            "created_at": datetime.now(),
+            "last_activity": datetime.now(),
+            "current_file": None,
+            "proposed_pipeline": None,
+            "state": "initial",  # initial, pipeline_proposed, executing, completed
+            "conversation_history": [],
+            "pipeline_executions": [],
+            "metadata": metadata or {},
+            "stats": {
+                "total_messages": 0,
+                "total_pipelines_executed": 0,
+                "total_tokens_used": 0
+            }
+        }
+        
+        if self.sessions_col is not None:
+            try:
+                self.sessions_col.insert_one(session_data)
+                print(f"✅ Session created in MongoDB: {session_id}")
+            except Exception as e:
+                print(f"⚠️  Failed to save session to MongoDB: {e}")
+        
+        return session_id
+    
+    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve session by ID
+        
+        Args:
+            session_id: Session identifier
+        
+        Returns:
+            Session data or None if not found
+        """
+        if self.sessions_col is None:
+            return None
+        
+        try:
+            session = self.sessions_col.find_one({"session_id": session_id})
+            if session:
+                # Update last activity
+                self.sessions_col.update_one(
+                    {"session_id": session_id},
+                    {"$set": {"last_activity": datetime.now()}}
+                )
+                # Remove MongoDB _id field
+                session.pop("_id", None)
+                return session
+            return None
+        except Exception as e:
+            print(f"⚠️  Error retrieving session: {e}")
+            return None
+    
+    def update_session(
+        self,
+        session_id: str,
+        updates: Dict[str, Any]
+    ) -> bool:
+        """
+        Update session data
+        
+        Args:
+            session_id: Session identifier
+            updates: Dictionary of fields to update
+        
+        Returns:
+            True if successful, False otherwise
+        """
+        if self.sessions_col is None:
+            return False
+        
+        try:
+            updates["last_activity"] = datetime.now()
+            
+            result = self.sessions_col.update_one(
+                {"session_id": session_id},
+                {"$set": updates}
+            )
+            
+            return result.modified_count > 0
+        except Exception as e:
+            print(f"⚠️  Error updating session: {e}")
+            return False
+    
+    def add_message(
+        self,
+        session_id: str,
+        role: str,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Add a message to conversation history
+        
+        Args:
+            session_id: Session identifier
+            role: Message role (user, assistant, system)
+            content: Message content
+            metadata: Additional message metadata
+        
+        Returns:
+            True if successful
+        """
+        if self.sessions_col is None:
+            return False
+        
+        try:
+            message = {
+                "role": role,
+                "content": content,
+                "timestamp": datetime.now(),
+                "metadata": metadata or {}
+            }
+            
+            self.sessions_col.update_one(
+                {"session_id": session_id},
+                {
+                    "$push": {"conversation_history": message},
+                    "$inc": {"stats.total_messages": 1},
+                    "$set": {"last_activity": datetime.now()}
+                }
+            )
+            
+            return True
+        except Exception as e:
+            print(f"⚠️  Error adding message: {e}")
+            return False
+    
+    def save_pipeline_execution(
+        self,
+        session_id: str,
+        pipeline: Dict[str, Any],
+        result: Dict[str, Any],
+        file_path: Optional[str] = None,
+        executor: str = "unknown"
+    ) -> bool:
+        """
+        Save pipeline execution to dedicated collection
+        
+        Args:
+            session_id: Session identifier
+            pipeline: Pipeline configuration
+            result: Execution result
+            file_path: File that was processed
+            executor: Which executor was used (bedrock, crewai, gemini)
+        
+        Returns:
+            True if successful
+        """
+        if self.pipelines_col is None:
+            return False
+        
+        try:
+            execution_data = {
+                "execution_id": str(uuid.uuid4()),
+                "session_id": session_id,
+                "pipeline_name": pipeline.get("pipeline_name"),
+                "pipeline_config": pipeline,
+                "result": result,
+                "file_path": file_path,
+                "executor": executor,
+                "executed_at": datetime.now(),
+                "duration_seconds": result.get("summary", {}).get("total_duration_seconds"),
+                "status": result.get("status", "unknown"),
+                "components_executed": len(pipeline.get("components", []))
+            }
+            
+            self.pipelines_col.insert_one(execution_data)
+            
+            # Update session stats
+            self.sessions_col.update_one(
+                {"session_id": session_id},
+                {
+                    "$inc": {"stats.total_pipelines_executed": 1},
+                    "$push": {"pipeline_executions": execution_data["execution_id"]}
+                }
+            )
+            
+            print(f"✅ Pipeline execution saved: {execution_data['execution_id']}")
+            return True
+            
+        except Exception as e:
+            print(f"⚠️  Error saving pipeline execution: {e}")
+            return False
+    
+    def get_session_history(
+        self,
+        session_id: str,
+        limit: int = 50
+    ) -> List[Dict[str, Any]]:
+        """
+        Get conversation history for a session
+        
+        Args:
+            session_id: Session identifier
+            limit: Maximum number of messages to return
+        
+        Returns:
+            List of messages
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return []
+        
+        history = session.get("conversation_history", [])
+        return history[-limit:] if len(history) > limit else history
+    
+    def get_pipeline_executions(
+        self,
+        session_id: Optional[str] = None,
+        limit: int = 10
+    ) -> List[Dict[str, Any]]:
+        """
+        Get pipeline execution history
+        
+        Args:
+            session_id: Optional session filter
+            limit: Maximum number of executions to return
+        
+        Returns:
+            List of pipeline executions
+        """
+        if self.pipelines_col is None:
+            return []
+        
+        try:
+            query = {"session_id": session_id} if session_id else {}
+            
+            executions = self.pipelines_col.find(query).sort("executed_at", -1).limit(limit)
+            
+            result = []
+            for exec_doc in executions:
+                exec_doc.pop("_id", None)
+                # Convert datetime to ISO string
+                if "executed_at" in exec_doc and isinstance(exec_doc["executed_at"], datetime):
+                    exec_doc["executed_at"] = exec_doc["executed_at"].isoformat()
+                result.append(exec_doc)
+            
+            return result
+            
+        except Exception as e:
+            print(f"⚠️  Error retrieving pipeline executions: {e}")
+            return []
+    
+    def cleanup_old_sessions(self, max_age_hours: int = 24) -> int:
+        """
+        Remove sessions older than max_age_hours
+        
+        Args:
+            max_age_hours: Maximum session age in hours
+        
+        Returns:
+            Number of sessions removed
+        """
+        if self.sessions_col is None:
+            return 0
+        
+        try:
+            cutoff = datetime.now() - timedelta(hours=max_age_hours)
+            
+            result = self.sessions_col.delete_many({
+                "last_activity": {"$lt": cutoff}
+            })
+            
+            count = result.deleted_count
+            if count > 0:
+                print(f"🧹 Cleaned up {count} old sessions")
+            
+            return count
+            
+        except Exception as e:
+            print(f"⚠️  Error cleaning up sessions: {e}")
+            return 0
+    
+    def get_session_stats(self, session_id: str) -> Dict[str, Any]:
+        """
+        Get statistics for a session
+        
+        Args:
+            session_id: Session identifier
+        
+        Returns:
+            Session statistics
+        """
+        session = self.get_session(session_id)
+        if not session:
+            return {}
+        
+        return {
+            "session_id": session_id,
+            "created_at": session.get("created_at"),
+            "last_activity": session.get("last_activity"),
+            "total_messages": session.get("stats", {}).get("total_messages", 0),
+            "total_pipelines_executed": session.get("stats", {}).get("total_pipelines_executed", 0),
+            "conversation_length": len(session.get("conversation_history", [])),
+            "state": session.get("state", "unknown")
+        }
+    
+    def close(self):
+        """Close MongoDB connection"""
+        if self.client:
+            self.client.close()
+            print("🔒 MongoDB connection closed")
+
+
+# Global session manager instance
+session_manager = SessionManager()
+
+
+if __name__ == "__main__":
+    # Test session manager
+    print("Testing Session Manager...")
+    
+    # Create session
+    sid = session_manager.create_session(user_id="test_user")
+    print(f"Created session: {sid}")
+    
+    # Add messages
+    session_manager.add_message(sid, "user", "Hello!")
+    session_manager.add_message(sid, "assistant", "Hi! How can I help?")
+    
+    # Get session
+    session = session_manager.get_session(sid)
+    print(f"Session data: {session}")
+    
+    # Get history
+    history = session_manager.get_session_history(sid)
+    print(f"History: {history}")
+    
+    # Get stats
+    stats = session_manager.get_session_stats(sid)
+    print(f"Stats: {stats}")