/ examples / persistence / sqlite_conversation_store.py
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!")