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