1  """
  2  MongoDB StateStore — Full Persistence Example
  3  
  4  Demonstrates:
  5    - Creating a MongoDB StateStore
  6    - Key-value operations (get/set/delete/exists)
  7    - Complex state persistence (agent metadata, usage tokens)
  8    - Direct MongoDB verification
  9    - Session resume after simulated restart
 10  
 11  Requirements:
 12      pip install praisonai pymongo
 13      Docker: MongoDB on localhost:27017 (no auth)
 14  
 15  Run:
 16      python mongodb_state_store.py
 17  """
 18  
 19  import uuid
 20  import pymongo
 21  from praisonai.persistence.state.mongodb import MongoDBStateStore
 22  
 23  MONGO_URL = "mongodb://localhost:27017"
 24  DATABASE = "praisonai_example"
 25  COLLECTION = f"state_{uuid.uuid4().hex[:8]}"
 26  
 27  print(f"MongoDB: {MONGO_URL}")
 28  print(f"Database: {DATABASE}, Collection: {COLLECTION}\n")
 29  
 30  store = MongoDBStateStore(url=MONGO_URL, database=DATABASE, collection=COLLECTION)
 31  
 32  # --- Phase 1: Basic key-value operations ---
 33  print("=== Phase 1: Basic Key-Value Operations ===")
 34  store.set("agent_id", "mongo_demo_agent_001")
 35  store.set("model", "gpt-4o-mini")
 36  store.set("counter", 42)
 37  
 38  print(f"  agent_id: {store.get('agent_id')}")
 39  print(f"  model: {store.get('model')}")
 40  print(f"  counter: {store.get('counter')}")
 41  print(f"  exists('agent_id'): {store.exists('agent_id')}")
 42  print(f"  exists('missing'): {store.exists('missing')}")
 43  
 44  assert store.get("agent_id") == "mongo_demo_agent_001"
 45  assert store.get("counter") == 42
 46  
 47  print()
 48  
 49  # --- Phase 2: Complex state persistence ---
 50  print("=== Phase 2: Complex State Persistence ===")
 51  managed_state = {
 52      "agent_id": "mongo_demo_agent_001",
 53      "agent_version": 4,
 54      "environment_id": "env_mongo_001",
 55      "total_input_tokens": 800,
 56      "total_output_tokens": 300,
 57      "compute_instance_id": "e2b_mongo_demo",
 58      "session_history": [
 59          {"id": "session_m1", "status": "completed"},
 60          {"id": "session_m2", "status": "idle"},
 61      ],
 62  }
 63  
 64  store.set("managed_state", managed_state)
 65  recovered = store.get("managed_state")
 66  
 67  print(f"  agent_id: {recovered['agent_id']}")
 68  print(f"  agent_version: {recovered['agent_version']}")
 69  print(f"  total_input_tokens: {recovered['total_input_tokens']}")
 70  print(f"  compute_instance_id: {recovered['compute_instance_id']}")
 71  print(f"  session_history count: {len(recovered['session_history'])}")
 72  
 73  assert recovered["agent_id"] == "mongo_demo_agent_001"
 74  assert recovered["total_input_tokens"] == 800
 75  
 76  print()
 77  
 78  # --- Phase 3: Direct MongoDB Verification ---
 79  print("=== Phase 3: Direct MongoDB Verification ===")
 80  client = pymongo.MongoClient(MONGO_URL)
 81  doc = client[DATABASE][COLLECTION].find_one({"_id": "managed_state"})
 82  print(f"  Raw doc _id: {doc['_id']}")
 83  print(f"  Raw doc value.agent_id: {doc['value']['agent_id']}")
 84  print(f"  Raw doc value.total_input_tokens: {doc['value']['total_input_tokens']}")
 85  assert doc["value"]["compute_instance_id"] == "e2b_mongo_demo"
 86  client.close()
 87  
 88  print()
 89  
 90  # --- Phase 4: Session Resume ---
 91  print("=== Phase 4: Session Resume (Simulating Restart) ===")
 92  store.close()
 93  
 94  store2 = MongoDBStateStore(url=MONGO_URL, database=DATABASE, collection=COLLECTION)
 95  
 96  recovered2 = store2.get("managed_state")
 97  assert recovered2 is not None, "State not found after restart!"
 98  print(f"  Recovered agent_id: {recovered2['agent_id']}")
 99  print(f"  Recovered tokens: in={recovered2['total_input_tokens']}, out={recovered2['total_output_tokens']}")
100  print(f"  Recovered compute: {recovered2['compute_instance_id']}")
101  
102  # Update
103  recovered2["total_input_tokens"] += 400
104  store2.set("managed_state", recovered2)
105  final = store2.get("managed_state")
106  print(f"  Updated input_tokens: {final['total_input_tokens']}")
107  assert final["total_input_tokens"] == 1200
108  
109  # Delete test
110  store2.delete("counter")
111  assert not store2.exists("counter")
112  print("  Delete verified: counter removed")
113  
114  store2.close()
115  
116  # --- Cleanup ---
117  print("\n=== Cleanup ===")
118  client = pymongo.MongoClient(MONGO_URL)
119  client[DATABASE].drop_collection(COLLECTION)
120  client.close()
121  print("  Collection dropped.")
122  
123  print("\n✅ MongoDB StateStore — All tests passed!")