Spaces:
Sleeping
Sleeping
| import uuid | |
| import logging | |
| # Configure logging | |
| logger = logging.getLogger(__name__) | |
| class SessionManager: | |
| """ | |
| Manages user sessions for processing and query operations | |
| """ | |
| def __init__(self): | |
| """Initialize the session manager with an empty sessions dictionary""" | |
| self.sessions = {} | |
| logger.info("SessionManager initialized") | |
| def create_session(self, initial_value="processing"): | |
| """Create a new session with a UUID and return the session ID""" | |
| session_id = str(uuid.uuid4()) | |
| self.sessions[session_id] = initial_value | |
| logger.info(f"Created session: {session_id} with initial value: {initial_value}") | |
| return session_id | |
| def update_session(self, session_id, value): | |
| """Update the session with the given ID to the provided value""" | |
| self.sessions[session_id] = value | |
| logger.info(f"Updated session: {session_id}") | |
| def session_exists(self, session_id): | |
| """Check if a session with the given ID exists""" | |
| return session_id in self.sessions | |
| def get_session(self, session_id): | |
| """Get the session data for the given ID""" | |
| return self.sessions.get(session_id) | |
| def get_sessions_summary(self): | |
| """Get a summary of all sessions""" | |
| session_info = {} | |
| for session_id, session_data in self.sessions.items(): | |
| if session_data == "processing": | |
| status = "processing" | |
| elif session_data == "failed": | |
| status = "failed" | |
| else: | |
| # It's a RAG pipeline object | |
| status = "ready" | |
| # Count documents in vector DB | |
| doc_count = 0 | |
| if hasattr(session_data, 'vector_db_retriever'): | |
| if hasattr(session_data.vector_db_retriever, 'documents'): | |
| doc_count = len(session_data.vector_db_retriever.documents) | |
| elif hasattr(session_data.vector_db_retriever, 'vectors'): | |
| doc_count = len(session_data.vector_db_retriever.vectors) | |
| session_info[session_id] = { | |
| "status": status | |
| } | |
| if status == "ready": | |
| session_info[session_id]["document_count"] = doc_count | |
| return { | |
| "session_count": len(self.sessions), | |
| "sessions": session_info, | |
| "memory_usage": { | |
| "session_count": len(self.sessions) | |
| } | |
| } |