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!")