sqlite_conversation_store.py
1 """ 2 SQLite ConversationStore ā Full Persistence Example 3 4 Demonstrates: 5 - Creating a SQLite ConversationStore 6 - Creating sessions and adding messages 7 - Verifying data is stored in the SQLite file 8 - Session resume: destroy agent, restore from DB, verify state 9 10 Requirements: 11 pip install praisonai 12 13 Run: 14 python sqlite_conversation_store.py 15 """ 16 17 import os 18 import sqlite3 19 import tempfile 20 from praisonai.persistence.conversation.sqlite import SQLiteConversationStore 21 from praisonai.persistence.conversation.base import ConversationSession, ConversationMessage 22 23 # --- Setup --- 24 db_path = os.path.join(tempfile.gettempdir(), "praison_sqlite_example.db") 25 print(f"SQLite DB path: {db_path}\n") 26 27 store = SQLiteConversationStore(path=db_path) 28 29 SESSION_ID = "sqlite-demo-session-001" 30 31 # --- Phase 1: Create session and add messages --- 32 print("=== Phase 1: Create Session & Add Messages ===") 33 session = ConversationSession( 34 session_id=SESSION_ID, 35 agent_id="sqlite_demo_agent", 36 name="SQLite Demo Session", 37 metadata={"agent_version": 3, "environment": "local"}, 38 ) 39 store.create_session(session) 40 41 messages = [ 42 ("user", "Hello! What is the capital of France?"), 43 ("assistant", "The capital of France is Paris."), 44 ("user", "What about Germany?"), 45 ("assistant", "The capital of Germany is Berlin."), 46 ] 47 48 for role, content in messages: 49 msg = ConversationMessage(session_id=SESSION_ID, role=role, content=content) 50 store.add_message(SESSION_ID, msg) 51 print(f" Added [{role}]: {content[:50]}") 52 53 print() 54 55 # --- Phase 2: Verify data in DB --- 56 print("=== Phase 2: Verify Data in SQLite ===") 57 retrieved_session = store.get_session(SESSION_ID) 58 print(f" Session ID: {retrieved_session.session_id}") 59 print(f" Agent ID: {retrieved_session.agent_id}") 60 print(f" Metadata: {retrieved_session.metadata}") 61 62 stored_messages = store.get_messages(SESSION_ID) 63 print(f" Messages stored: {len(stored_messages)}") 64 for msg in stored_messages: 65 print(f" [{msg.role}] {msg.content[:60]}") 66 67 print() 68 69 # --- Phase 3: Session Resume --- 70 print("=== Phase 3: Session Resume (Simulating Restart) ===") 71 store.close() 72 73 # Reopen store from same file (simulating process restart) 74 store2 = SQLiteConversationStore(path=db_path) 75 76 # Verify session still exists 77 resumed_session = store2.get_session(SESSION_ID) 78 assert resumed_session is not None, "Session not found after restart!" 79 print(f" Resumed session: {resumed_session.session_id}") 80 print(f" Agent ID: {resumed_session.agent_id}") 81 print(f" Metadata preserved: {resumed_session.metadata}") 82 83 # Verify messages 84 resumed_messages = store2.get_messages(SESSION_ID) 85 print(f" Messages recovered: {len(resumed_messages)}") 86 assert len(resumed_messages) == 4, f"Expected 4 messages, got {len(resumed_messages)}" 87 88 # Continue the conversation 89 new_msg = ConversationMessage(session_id=SESSION_ID, role="user", content="What about Spain?") 90 store2.add_message(SESSION_ID, new_msg) 91 new_msg2 = ConversationMessage(session_id=SESSION_ID, role="assistant", content="The capital of Spain is Madrid.") 92 store2.add_message(SESSION_ID, new_msg2) 93 94 final_messages = store2.get_messages(SESSION_ID) 95 print(f" Messages after resume: {len(final_messages)}") 96 assert len(final_messages) == 6 97 98 store2.close() 99 100 # --- Phase 4: Direct SQLite verification --- 101 print("\n=== Phase 4: Direct SQLite Verification ===") 102 conn = sqlite3.connect(db_path) 103 cur = conn.cursor() 104 cur.execute("SELECT COUNT(*) FROM praison_messages WHERE session_id = ?", (SESSION_ID,)) 105 count = cur.fetchone()[0] 106 print(f" Raw SQL message count: {count}") 107 assert count == 6 108 conn.close() 109 110 # Cleanup 111 os.remove(db_path) 112 print("\nā SQLite ConversationStore ā All tests passed!")