/ examples / persistence / json_file_session_store.py
json_file_session_store.py
  1  """
  2  JSON File DefaultSessionStore — Full Persistence Example
  3  
  4  Demonstrates:
  5    - Creating a DefaultSessionStore (JSON file-backed)
  6    - Adding messages and metadata
  7    - Verifying data on disk
  8    - Session resume after simulated restart
  9    - Full managed agent state roundtrip (7 data categories)
 10  
 11  Requirements:
 12      pip install praisonaiagents
 13  
 14  Run:
 15      python json_file_session_store.py
 16  """
 17  
 18  import json
 19  import os
 20  import shutil
 21  import tempfile
 22  import uuid
 23  from praisonaiagents.session.store import DefaultSessionStore
 24  
 25  # --- Setup ---
 26  session_dir = os.path.join(tempfile.gettempdir(), f"praison_json_example_{uuid.uuid4().hex[:6]}")
 27  os.makedirs(session_dir, exist_ok=True)
 28  print(f"Session directory: {session_dir}\n")
 29  
 30  store = DefaultSessionStore(session_dir=session_dir)
 31  SESSION_ID = f"json-demo-{uuid.uuid4().hex[:8]}"
 32  
 33  # --- Phase 1: Add messages ---
 34  print("=== Phase 1: Add Messages ===")
 35  messages = [
 36      ("user", "What is machine learning?"),
 37      ("assistant", "Machine learning is a subset of AI that enables systems to learn from data."),
 38      ("user", "Give me an example."),
 39      ("assistant", "Spam filters use ML to classify emails as spam or not spam."),
 40  ]
 41  
 42  for role, content in messages:
 43      store.add_message(SESSION_ID, role, content)
 44      print(f"  Added [{role}]: {content[:60]}")
 45  
 46  print()
 47  
 48  # --- Phase 2: Store metadata (simulating managed agent state) ---
 49  print("=== Phase 2: Store Metadata ===")
 50  session = store.get_session(SESSION_ID)
 51  session.metadata = {
 52      "agent_id": "json_demo_agent_001",
 53      "agent_version": 2,
 54      "environment_id": "env_local_json",
 55      "total_input_tokens": 450,
 56      "total_output_tokens": 180,
 57      "compute_instance_id": "local_json_demo",
 58      "session_history": [
 59          {"id": SESSION_ID, "status": "idle", "title": "ML chat"},
 60      ],
 61  }
 62  store._save_session(session)
 63  
 64  print(f"  agent_id: {session.metadata['agent_id']}")
 65  print(f"  tokens: in={session.metadata['total_input_tokens']}, out={session.metadata['total_output_tokens']}")
 66  print(f"  compute: {session.metadata['compute_instance_id']}")
 67  
 68  print()
 69  
 70  # --- Phase 3: Verify on disk ---
 71  print("=== Phase 3: Verify JSON File on Disk ===")
 72  json_file = os.path.join(session_dir, f"{SESSION_ID}.json")
 73  assert os.path.exists(json_file), f"JSON file not found: {json_file}"
 74  
 75  with open(json_file) as f:
 76      data = json.load(f)
 77  
 78  print(f"  File: {json_file}")
 79  print(f"  Session ID: {data['session_id']}")
 80  print(f"  Messages: {len(data.get('messages', []))}")
 81  print(f"  Metadata agent_id: {data['metadata']['agent_id']}")
 82  print(f"  Metadata tokens_in: {data['metadata']['total_input_tokens']}")
 83  
 84  assert data["metadata"]["agent_id"] == "json_demo_agent_001"
 85  assert data["metadata"]["total_input_tokens"] == 450
 86  assert len(data["messages"]) == 4
 87  
 88  print()
 89  
 90  # --- Phase 4: Session Resume ---
 91  print("=== Phase 4: Session Resume (Simulating Restart) ===")
 92  del store
 93  
 94  store2 = DefaultSessionStore(session_dir=session_dir)
 95  
 96  # Verify session exists
 97  assert store2.session_exists(SESSION_ID), "Session not found after restart!"
 98  
 99  # Recover history
100  history = store2.get_chat_history(SESSION_ID)
101  print(f"  Recovered messages: {len(history)}")
102  assert len(history) == 4
103  
104  # Recover metadata
105  session2 = store2.get_session(SESSION_ID)
106  meta = session2.metadata
107  print(f"  Recovered agent_id: {meta['agent_id']}")
108  print(f"  Recovered tokens: in={meta['total_input_tokens']}, out={meta['total_output_tokens']}")
109  print(f"  Recovered compute: {meta['compute_instance_id']}")
110  print(f"  Recovered session_history: {len(meta['session_history'])} entries")
111  
112  assert meta["agent_id"] == "json_demo_agent_001"
113  assert meta["total_input_tokens"] == 450
114  assert meta["compute_instance_id"] == "local_json_demo"
115  
116  # Continue conversation
117  store2.add_message(SESSION_ID, "user", "What about deep learning?")
118  store2.add_message(SESSION_ID, "assistant", "Deep learning uses neural networks with many layers.")
119  final_history = store2.get_chat_history(SESSION_ID)
120  print(f"  Messages after resume: {len(final_history)}")
121  assert len(final_history) == 6
122  
123  # --- Cleanup ---
124  print("\n=== Cleanup ===")
125  shutil.rmtree(session_dir)
126  print("  Session directory removed.")
127  
128  print("\nāœ… JSON File DefaultSessionStore — All tests passed!")