diff --git a/examples/agent_save_load_full.py b/examples/agent_save_load_full.py new file mode 100644 index 00000000..3e9bef83 --- /dev/null +++ b/examples/agent_save_load_full.py @@ -0,0 +1,121 @@ +""" +Example: Fully Save and Load an Agent (with Conversation History) + +This demonstrates how to: + 1. Auto-save conversation messages to JSON + 2. Save the full Agent state + 3. Load both the Agent state and the conversation back into a fresh Agent +""" + +import os +from swarms.structs.agent import Agent + +# Helper to safely print type or None for agent properties +def print_agent_properties(agent, label): + print(f"\n--- {label} ---") + for prop in ["tokenizer", "long_term_memory", "logger_handler", "agent_output", "executor"]: + value = getattr(agent, prop, None) + print(f"{prop}: {type(value)}") + +# Helper to extract the conversation history list +def get_conversation_history(agent): + conv = getattr(agent, "conversation", None) or getattr(agent, "short_memory", None) + return getattr(conv, "conversation_history", None) + +# Robust helper to reload conversation from JSON into the correct attribute +def reload_conversation_from_json(agent, filepath): + conv = getattr(agent, "conversation", None) or getattr(agent, "short_memory", None) + if conv and hasattr(conv, "load_from_json"): + conv.load_from_json(filepath) + +# --- 1. Setup: Create and configure an agent with auto-save conversation --- +agent = Agent( + agent_name="test", + user_name="test_user", + system_prompt="This is a test agent", + max_loops=1, + context_length=200000, + autosave=True, + verbose=True, + artifacts_on=True, + artifacts_output_path="test", + artifacts_file_extension=".txt", + conversation_kwargs={ + "auto_save": True, + "save_as_json_bool": True, + "save_filepath": "test_conversation_history.json" + } +) + +# --- 2. Interact to populate conversation --- +agent.run(task="hello") +agent.run(task="What is your purpose?") +agent.run(task="Tell me a joke.") +agent.run(task="Summarize our conversation so far.") + +# --- 3. Inspect before saving --- +print_agent_properties(agent, "BEFORE SAVE") +print("\nConversation history BEFORE SAVE:", get_conversation_history(agent)) + +# --- 4. Save the agent state (conversation JSON was auto-saved under workspace) --- +state_path = os.path.join(agent.workspace_dir, "test_state.json") +agent.save(state_path) + +# --- 5. Ensure the conversation JSON file is saved and print its path and contents --- +json_path = os.path.join(agent.workspace_dir, "test_conversation_history.json") +if hasattr(agent, "short_memory") and hasattr(agent.short_memory, "save_as_json"): + agent.short_memory.save_as_json(json_path) + +if os.path.exists(json_path): + print(f"\n[CHECK] Conversation JSON file found: {json_path}") + with open(json_path, "r") as f: + json_data = f.read() + print("[CHECK] JSON file contents:\n", json_data) +else: + print(f"[WARN] Conversation JSON file not found: {json_path}") + +# --- 6. Simulate fresh environment --- +del agent + +# --- 7. Load: Restore the agent configuration --- +agent2 = Agent(agent_name="test") +agent2.load(state_path) + +# --- 8. Load: Restore the conversation history from the workspace directory into a new Conversation object --- +from swarms.structs.conversation import Conversation +conversation_loaded = Conversation() +if os.path.exists(json_path): + conversation_loaded.load_from_json(json_path) + print("\n[CHECK] Loaded conversation from JSON into new Conversation object:") + print(conversation_loaded.conversation_history) +else: + print(f"[WARN] Conversation JSON file not found for loading: {json_path}") + +# --- 9. Assign loaded conversation to agent2 and check --- +agent2.short_memory = conversation_loaded +print("\n[CHECK] Agent2 conversation history after assigning loaded conversation:", get_conversation_history(agent2)) + +# --- 10. Inspect after loading --- +print_agent_properties(agent2, "AFTER LOAD") +print("\nConversation history AFTER LOAD:", get_conversation_history(agent2)) + +# --- 11. Confirm the agent can continue --- +result = agent2.run(task="What is 2+2?") +print("\nAgent2 run result:", result) + +# --- 12. Cleanup test files --- +print(f"\n[INFO] Test complete. Conversation JSON and agent state files are available for inspection:") +print(f" Conversation JSON: {json_path}") +print(f" Agent state: {state_path}") +print("You can open and inspect these files to verify the agent's memory persistence.") +# Do NOT delete files automatically +# for path in (state_path, json_path): +# try: +# os.remove(path) +# except OSError: +# pass + +# --- 13. Test if agent2 remembers the previous conversation --- +print("\n[TEST] Checking if agent2 remembers the previous conversation after reload...") +probe = agent2.run(task="What did I ask you to do earlier?") +print("\nAgent2 memory probe result:", probe) diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py index 74c1ccab..988e262b 100644 --- a/swarms/structs/agent.py +++ b/swarms/structs/agent.py @@ -67,9 +67,11 @@ from swarms.utils.pdf_to_text import pdf_to_text from swarms.utils.str_to_dict import str_to_dict from swarms.prompts.react_base_prompt import REACT_SYS_PROMPT from swarms.prompts.max_loop_prompt import generate_reasoning_prompt +from swarms.structs.agent_non_serializable import restore_non_serializable_properties from swarms.prompts.safety_prompt import SAFETY_PROMPT + # Utils # Custom stopping condition def stop_when_repeats(response: str) -> bool: @@ -1725,6 +1727,9 @@ class Agent: # Reinitialize any necessary runtime components self._reinitialize_after_load() + # Restore non-serializable properties (tokenizer, long_term_memory, logger_handler, agent_output, executor) + self.restore_non_serializable_properties() + if self.verbose: self._log_loaded_state_info(resolved_path) @@ -2781,3 +2786,25 @@ class Agent: role="Output Cleaner", content=response, ) + + def restore_non_serializable_properties(self): + """ + Restore non-serializable properties for the Agent instance after loading. + This should be called after loading agent state from disk. + """ + restore_non_serializable_properties(self) + + # Custom serialization for non-serializable properties + def __getstate__(self): + state = self.__dict__.copy() + # Remove non-serializable properties + for prop in ["tokenizer", "long_term_memory", "logger_handler", "agent_output", "executor"]: + if prop in state: + state[prop] = None # Or a serializable placeholder if needed + return state + + def __setstate__(self, state): + self.__dict__.update(state) + # Restore non-serializable properties after loading + if hasattr(self, 'restore_non_serializable_properties'): + self.restore_non_serializable_properties() diff --git a/swarms/structs/agent_non_serializable.py b/swarms/structs/agent_non_serializable.py new file mode 100644 index 00000000..1342899f --- /dev/null +++ b/swarms/structs/agent_non_serializable.py @@ -0,0 +1,62 @@ +""" +Non-Serializable Properties Handler for Agent + +This module provides helper functions to save and restore non-serializable properties +(tokenizer, long_term_memory, logger_handler, agent_output, executor) for the Agent class. + +Usage: + from swarms.structs.agent_non_serializable import restore_non_serializable_properties + restore_non_serializable_properties(agent) +""" + +from concurrent.futures import ThreadPoolExecutor +import logging + +# Dummy/placeholder for long_term_memory and agent_output restoration +class DummyLongTermMemory: + def __init__(self): + self.memory = [] + def query(self, *args, **kwargs): + # Return an empty list or a default value to avoid errors + return [] + def save(self, path): + # Optionally implement a no-op save for compatibility + pass + +class DummyAgentOutput: + def __init__(self): + self.output = None + +def restore_non_serializable_properties(agent): + """ + Restore non-serializable properties for the Agent instance after loading. + This should be called after loading agent state from disk. + """ + # Restore tokenizer using LiteLLM if available + agent.tokenizer = None + try: + from swarms.utils.litellm_tokenizer import count_tokens + agent.tokenizer = count_tokens # Assign the function as a tokenizer interface + except Exception: + agent.tokenizer = None + + # Restore long_term_memory (dummy for demo, replace with real backend as needed) + if getattr(agent, "long_term_memory", None) is None or not hasattr(agent.long_term_memory, "query"): + agent.long_term_memory = DummyLongTermMemory() + + # Restore logger_handler + try: + agent.logger_handler = logging.StreamHandler() + except Exception: + agent.logger_handler = None + + # Restore agent_output (dummy for demo, replace with real backend as needed) + agent.agent_output = DummyAgentOutput() + + # Restore executor + try: + agent.executor = ThreadPoolExecutor() + except Exception: + agent.executor = None + + return agent