Spaces:
Sleeping
Sleeping
File size: 2,579 Bytes
209e402 |
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 |
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)
}
} |