import sys from typing import Dict, List, Tuple from constants import MAX_FILE_SIZES_PER_SESSION class SessionDocumentStore: def __init__(self) -> None: # Stores, for each session, the files' content and name # session_id -> {file_name -> (file_text, size_in_bytes)} self.session_document_map: Dict[str, Dict[str, Tuple[str, int]]] = dict() def create_document(self, session_id: str, file_text: str, file_name: str): text_size = sys.getsizeof(file_text) if session_id not in self.session_document_map: if text_size > MAX_FILE_SIZES_PER_SESSION: return False self.session_document_map[session_id] = dict() self.session_document_map[session_id][file_name] = (file_text, text_size) return True current_total_file_size = sum( file_text_size[1] for file_text_size in self.session_document_map[session_id].values() ) if current_total_file_size + text_size > MAX_FILE_SIZES_PER_SESSION: return False self.session_document_map[session_id][file_name] = (file_text, text_size) return True def get_document_contents(self, session_id: str) -> List[str] | None: document_map = self.session_document_map.get(session_id) if document_map is None: return None document_contents = [ file_text_size[0] for file_text_size in document_map.values() ] if len(document_contents) == 0: return None return document_contents def delete_document(self, session_id: str, file_name: str) -> bool: """Deletes a document with the given name. If the session no longer has documents after the deletion, the session is also deleted and the function returns True.""" document_map = self.session_document_map.get(session_id) if document_map is None: return False if file_name in document_map: del document_map[file_name] if len(document_map) == 0: self.delete_session_documents(session_id) return True return False def delete_session_documents(self, session_id: str) -> bool: return self.session_document_map.pop(session_id, None) is not None