From bef75c9f2e05cb3b4a8089ff2874da0e99c5eea2 Mon Sep 17 00:00:00 2001 From: ascender1729 Date: Fri, 9 May 2025 21:45:42 +0530 Subject: [PATCH] fix(agent): Implement logger handler serialization - Add _serialize_logger_handler and _deserialize_logger_handler methods - Update save/load methods to handle logger serialization - Add real-world usage test - Fix file path handling for cross-platform compatibility Resolves #640 (logger handler component) --- swarms/structs/agent.py | 150 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py index f37f5d61..f2bd913f 100644 --- a/swarms/structs/agent.py +++ b/swarms/structs/agent.py @@ -2803,3 +2803,153 @@ class Agent: role="Output Cleaner", content=response, ) + + def _serialize_logger_handler(self) -> dict: + """Serialize logger handler configuration.""" + if not hasattr(self, 'logger_handler') or self.logger_handler is None: + return None + + handler_config = { + 'type': 'TextIOWrapper', + 'mode': self.logger_handler.mode, + 'name': self.logger_handler.name if hasattr(self.logger_handler, 'name') else None, + 'encoding': self.logger_handler.encoding if hasattr(self.logger_handler, 'encoding') else None + } + return handler_config + + def _deserialize_logger_handler(self, config: dict) -> None: + """Recreate logger handler from configuration.""" + if not config: + self.logger_handler = None + return + + try: + if config['type'] == 'TextIOWrapper': + if config['name']: + self.logger_handler = open( + config['name'], + mode=config.get('mode', 'a'), + encoding=config.get('encoding', 'utf-8') + ) + except Exception as e: + logger.warning(f"Failed to recreate logger handler: {e}") + self.logger_handler = None + + def _serialize_short_memory(self) -> dict: + """Serialize short memory into a dictionary format.""" + if hasattr(self, 'short_memory') and self.short_memory is not None: + return { + 'type': 'Conversation', + 'messages': self.short_memory.messages if hasattr(self.short_memory, 'messages') else [], + 'system_prompt': self.short_memory.system_prompt if hasattr(self.short_memory, 'system_prompt') else None, + 'time_enabled': self.short_memory.time_enabled if hasattr(self.short_memory, 'time_enabled') else False, + 'user': self.short_memory.user if hasattr(self.short_memory, 'user') else "Human:", + 'rules': self.short_memory.rules if hasattr(self.short_memory, 'rules') else None + } + return None + + def _deserialize_short_memory(self, config: dict) -> None: + """Recreate short memory from configuration.""" + if not config: + self.short_memory = self.short_memory_init() + return + + try: + if config['type'] == 'Conversation': + from swarms.structs.conversation import Conversation + self.short_memory = Conversation( + system_prompt=config.get('system_prompt'), + time_enabled=config.get('time_enabled', False), + user=config.get('user', "Human:"), + rules=config.get('rules') + ) + # Restore messages + if 'messages' in config: + self.short_memory.messages = config['messages'] + except Exception as e: + logger.warning(f"Failed to recreate short memory: {e}") + self.short_memory = self.short_memory_init() + + def save(self, file_path: str = None) -> None: + """Save agent state to file.""" + try: + # Get the full path for saving + full_path = file_path or self.saved_state_path + if not full_path: + raise ValueError("No file path provided for saving state") + + # Create temporary file for atomic save + temp_path = f"{full_path}.tmp" + backup_path = f"{full_path}.bak" + + # Prepare state dictionary + state_dict = self.__dict__.copy() + + # Handle special serialization for non-serializable objects + state_dict['logger_handler'] = self._serialize_logger_handler() + state_dict['short_memory'] = self._serialize_short_memory() + + # Remove other non-serializable objects + non_serializable = ['llm', 'tokenizer', 'long_term_memory', 'agent_output', 'executor'] + for key in non_serializable: + if key in state_dict: + if state_dict[key] is not None: + state_dict[key] = f"" + else: + state_dict[key] = None + + # Save to temporary file + with open(temp_path, 'w') as f: + json.dump(state_dict, f, indent=2, default=str) + + # If current file exists, create backup + if os.path.exists(full_path): + os.replace(full_path, backup_path) + + # Move temporary file to final location + os.replace(temp_path, full_path) + + # Clean up old backup if everything succeeded + if os.path.exists(backup_path): + try: + os.remove(backup_path) + except Exception as e: + logger.warning(f"Could not remove backup file: {e}") + + logger.info(f"Successfully saved agent state to: {full_path}") + + except Exception as e: + logger.error(f"Error saving agent state: {e}") + raise + + def load(self, file_path: str = None) -> None: + """Load agent state from file.""" + try: + # Get the full path for loading + full_path = file_path or self.saved_state_path + if not full_path: + raise ValueError("No file path provided for loading state") + + # Load state from file + with open(full_path, 'r') as f: + state_dict = json.load(f) + + # Handle special deserialization first + logger_config = state_dict.pop('logger_handler', None) + short_memory_config = state_dict.pop('short_memory', None) + + self._deserialize_logger_handler(logger_config) + self._deserialize_short_memory(short_memory_config) + + # Update remaining agent attributes + for key, value in state_dict.items(): + # Skip non-serializable objects that were saved as strings + if isinstance(value, str) and value.startswith('