/ examples / managed-agents / persistence / postgres_managed.py
postgres_managed.py
 1  #!/usr/bin/env python3
 2  """
 3  ManagedAgent + PostgreSQL Persistence — Real conversation with session resume.
 4  
 5  Prerequisites:
 6      pip install praisonai praisonaiagents psycopg2-binary
 7      export OPENAI_API_KEY="sk-..."
 8      # PostgreSQL running on localhost:5432
 9  """
10  
11  import os
12  import sys
13  
14  import psycopg2
15  
16  if not os.getenv("OPENAI_API_KEY"):
17      sys.exit("ERROR: OPENAI_API_KEY not set.")
18  
19  from praisonai import ManagedAgent, LocalManagedConfig, DB
20  from praisonaiagents import Agent
21  
22  PG_URL = os.getenv("PG_URL", "postgresql://postgres:postgres@localhost:5432/postgres")
23  print(f"ManagedAgent + PostgreSQL Persistence\nDB: {PG_URL}\n")
24  
25  # ── Phase 1: Create agent, teach facts ──
26  print("=== Phase 1: First Session ===")
27  
28  db = DB(database_url=PG_URL)
29  managed = ManagedAgent(
30      provider="local", db=db,
31      config=LocalManagedConfig(
32          model="gpt-4o-mini", name="PG Memory Agent",
33          system="You are a helpful assistant. Remember all facts the user tells you.",
34      ),
35  )
36  agent = Agent(name="User", backend=managed)
37  
38  result1 = agent.run("Remember: My favourite language is Rust, 3 years experience. Confirm.")
39  print(f"Agent: {result1[:200]}...")
40  print(f"Session ID: {managed.session_id}")
41  
42  result2 = agent.run("Also remember: I'm building a web crawler with the reqwest crate. Confirm.")
43  print(f"Agent: {result2[:200]}...")
44  
45  # ── Phase 2: Direct PostgreSQL Verification ──
46  print("\n=== Phase 2: PostgreSQL Verification ===")
47  conn = psycopg2.connect(PG_URL)
48  cursor = conn.cursor()
49  cursor.execute("""
50      SELECT table_name FROM information_schema.tables
51      WHERE table_schema = 'public' AND table_name LIKE '%message%'
52  """)
53  msg_tables = [r[0] for r in cursor.fetchall()]
54  print(f"Message tables: {msg_tables}")
55  for t in msg_tables:
56      cursor.execute(f"SELECT COUNT(*) FROM {t}")
57      print(f"  {t}: {cursor.fetchone()[0]} rows")
58  conn.close()
59  
60  # ── Phase 3: Destroy instance ──
61  print("\n=== Phase 3: Instance Goes Idle ===")
62  saved_ids = managed.save_ids()
63  print(f"Saved IDs: {saved_ids}")
64  del agent, managed, db
65  print("Agent destroyed.\n")
66  
67  # ── Phase 4: Resume — no config needed ──
68  print("=== Phase 4: Resume Session ===")
69  db2 = DB(database_url=PG_URL)
70  managed2 = ManagedAgent(provider="local", db=db2)
71  managed2.resume_session(saved_ids["session_id"])
72  print(f"Resumed session: {managed2.session_id}")
73  
74  agent2 = Agent(name="User", backend=managed2)
75  result3 = agent2.run("What is my favourite language, experience, and what am I building?")
76  print(f"Agent: {result3[:300]}...")
77  
78  # ── Phase 5: Validate ──
79  print("\n=== Phase 5: Validation ===")
80  r = result3.lower()
81  checks = {
82      "Remembers 'Rust'": "rust" in r,
83      "Remembers '3 years'": "3" in r and "year" in r,
84      "Remembers 'crawler'": "crawler" in r or "crawl" in r,
85      "Remembers 'reqwest'": "reqwest" in r,
86      "Session ID continuity": managed2.session_id == saved_ids["session_id"],
87  }
88  all_passed = True
89  for check, ok in checks.items():
90      print(f"  [{'PASS' if ok else 'FAIL'}] {check}")
91      if not ok:
92          all_passed = False
93  
94  print(f"\n{'PASS' if all_passed else 'FAIL'}: PostgreSQL ManagedAgent persistence test")