From 61cb4aab15e39f9d7eff7b4028d856428ff7fa9c Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Mon, 26 May 2025 13:31:58 +0000
Subject: [PATCH] WIP: robust agent save/load example for conversation autosave
 and reload

---
 examples/agent_save_load_full.py | 95 ++++++++++++++++++++++++--------
 1 file changed, 72 insertions(+), 23 deletions(-)

diff --git a/examples/agent_save_load_full.py b/examples/agent_save_load_full.py
index 758423f7..96860cd5 100644
--- a/examples/agent_save_load_full.py
+++ b/examples/agent_save_load_full.py
@@ -1,14 +1,14 @@
 """
-Example: Fully Save and Load an Agent (Issue #640)
+Example: Fully Save and Load an Agent (with Conversation History)
 
-This example demonstrates how to save and load an Agent instance such that all non-serializable properties
-(tokenizer, long_term_memory, logger_handler, agent_output, executor) are restored after loading.
-
-This is a user-facing, production-grade demonstration for swarms.
+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
 """
 
-from swarms.structs.agent import 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):
@@ -17,7 +17,18 @@ def print_agent_properties(agent, label):
         value = getattr(agent, prop, None)
         print(f"{prop}: {type(value)}")
 
-# --- Setup: Create and configure an agent ---
+# 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",
@@ -29,34 +40,72 @@ agent = Agent(
     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"
+    }
 )
 
-# Optionally, interact with the agent to populate state
+# --- 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.")
 
-# Print non-serializable properties BEFORE saving
+# --- 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)
 
-# Save the agent state
-save_path = os.path.join(agent.workspace_dir, "test_state.json")
-agent.save(save_path)
+# --- 5. Check that the conversation JSON file exists and print its contents ---
+json_path = os.path.join(agent.workspace_dir, "test_conversation_history.json")
+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}")
 
-# Delete the agent instance to simulate a fresh load
+# --- 6. Simulate fresh environment ---
 del agent
 
-# --- Load: Restore the agent from file ---
-agent2 = Agent(agent_name="test")  # Minimal init, will be overwritten by load
-agent2.load(save_path)
+# --- 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 ---
+if hasattr(agent2, "conversation"):
+    agent2.conversation = conversation_loaded
+elif hasattr(agent2, "short_memory"):
+    agent2.short_memory = conversation_loaded
+print("\n[CHECK] Agent2 conversation history after assigning loaded conversation:", get_conversation_history(agent2))
 
-# Print non-serializable properties AFTER loading
+# --- 10. Inspect after loading ---
 print_agent_properties(agent2, "AFTER LOAD")
+print("\nConversation history AFTER LOAD:", get_conversation_history(agent2))
 
-# Confirm agent2 can still run tasks and autosave
+# --- 11. Confirm the agent can continue ---
 result = agent2.run(task="What is 2+2?")
 print("\nAgent2 run result:", result)
 
-# Clean up test file
-try:
-    os.remove(save_path)
-except Exception:
-    pass
+# --- 12. Cleanup test files ---
+for path in (state_path, json_path):
+    try:
+        os.remove(path)
+    except OSError:
+        pass