import os import hashlib import base64 from typing import Optional, Dict, ClassVar from loguru import logger from pathlib import Path from autogen_ext.experimental.task_centric_memory import ( MemoryController, MemoryControllerConfig, MemoryBankConfig, ) from autogen_ext.experimental.task_centric_memory.utils import PageLogger from autogen_core.models import ChatCompletionClient MEMORY_SUBDIR = "memory_bank" LOG_SUBDIR = "pagelogs" class MemoryControllerProvider: """Singleton provider for memory controller instances""" _instance: ClassVar[Optional["MemoryControllerProvider"]] = None _memory_controllers: Dict[str, MemoryController] = {} _internal_workspace_root: Optional[Path] = None _external_workspace_root: Optional[Path] = None _inside_docker: bool = False def __new__( cls, internal_workspace_root: Optional[Path] = None, external_workspace_root: Optional[Path] = None, inside_docker: bool = False, ): if cls._instance is None: cls._instance = super(MemoryControllerProvider, cls).__new__(cls) cls._instance._memory_controllers = {} cls._instance._internal_workspace_root = internal_workspace_root cls._instance._external_workspace_root = external_workspace_root cls._instance._inside_docker = inside_docker return cls._instance def __init__( self, internal_workspace_root: Optional[Path] = None, external_workspace_root: Optional[Path] = None, inside_docker: bool = False, ): """ Initialize the memory controller provider with paths Args: internal_workspace_root (Path, optional): Path to workspace root inside docker external_workspace_root (Path, optional): Path to workspace root on host inside_docker (bool, optional): Whether code is running inside Docker. Default: False """ if internal_workspace_root is not None: self._internal_workspace_root = internal_workspace_root if external_workspace_root is not None: self._external_workspace_root = external_workspace_root self._inside_docker = inside_docker root = ( self._internal_workspace_root if self._inside_docker else self._external_workspace_root ) if isinstance(root, Path): memory_dir = root / MEMORY_SUBDIR log_dir = root / LOG_SUBDIR memory_dir.mkdir(exist_ok=True, parents=True) log_dir.mkdir(exist_ok=True, parents=True) try: self.validate_path_safety(memory_dir, root) self.validate_path_safety(log_dir, root) except ValueError as e: logger.error( f"Path safety validation failed during initialization: {e}" ) raise ValueError(f"Security violation: {e}") else: logger.warning("Invalid root path: root must be a Path object") @staticmethod def get_safe_key(memory_controller_key: str) -> str: """Convert a user ID into a filesystem-safe key""" hash_obj = hashlib.sha256(memory_controller_key.encode("utf-8")) safe_id = ( base64.urlsafe_b64encode(hash_obj.digest()[:16]).decode("utf-8").rstrip("=") ) return safe_id @staticmethod def validate_path_safety(path: Path, base_dir: Path) -> bool: """ Validate that a path is within the expected base directory. Helps prevent directory traversal attacks. Args: path (Path): The path to validate base_dir (Path): The base directory that should contain the path Returns: bool: True if path is safe Raises: ValueError: If the path is outside the base directory """ resolved_path = os.path.realpath(path) resolved_base = os.path.realpath(base_dir) if not resolved_path.startswith(resolved_base): raise ValueError(f"Path validation failed: {path} is outside of {base_dir}") return True def get_path(self, subdir: str, safe_key: str) -> Optional[Path]: """Get path based on current context (Docker or host)""" root = ( self._internal_workspace_root if self._inside_docker else self._external_workspace_root ) if isinstance(root, Path): full_path = root / subdir / safe_key full_path.mkdir(exist_ok=True, parents=True) try: self.validate_path_safety(full_path, root / subdir) return full_path except ValueError: raise ValueError("Invalid memory controller key") logger.warning( "Memory controller provider root path is not a valid Path object" ) return None def get_memory_controller( self, memory_controller_key: str, client: ChatCompletionClient, reset: bool = False, ) -> MemoryController: """Get or create a memory controller for the specified user""" safe_key = self.get_safe_key(memory_controller_key) if safe_key in self._memory_controllers and not reset: return self._memory_controllers[safe_key] memory_path = self.get_path(MEMORY_SUBDIR, safe_key) log_path = self.get_path(LOG_SUBDIR, safe_key) try: page_logger = PageLogger(config={"level": "INFO", "path": str(log_path)}) memory_bank_config = MemoryBankConfig( path=str(memory_path), relevance_conversion_threshold=float( os.environ.get("MEMORY_RELEVANCE_THRESHOLD", "1.7") ), ) memory_controller_config = MemoryControllerConfig( generalize_task=False, revise_generalized_task=False, generate_topics=False, validate_memos=True, max_memos_to_retrieve=1, MemoryBank=memory_bank_config, ) memory_controller = MemoryController( reset=reset, client=client, logger=page_logger, config=memory_controller_config, ) self._memory_controllers[safe_key] = memory_controller return memory_controller except Exception as e: logger.error(f"Error creating memory controller: {e}") raise def close_memory_controller(self, memory_controller_key: str) -> None: """Close a memory controller and clean up resources""" safe_key = self.get_safe_key(memory_controller_key) if safe_key in self._memory_controllers: logger.info(f"Closing memory controller (safe key: {safe_key})") del self._memory_controllers[safe_key] def close_all_memory_controllers(self) -> None: """Close all memory controllers""" logger.info( f"Closing all memory controllers ({len(self._memory_controllers)} total)" ) self._memory_controllers.clear()