File size: 4,583 Bytes
bb17288 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import os
import psycopg2
import uuid
import datetime
import logging
import json
logger = logging.getLogger(__name__)
class PostgresMemoryManager:
def __init__(self):
self.db_url = os.getenv("NEON_DATABASE_URL")
self.conn = None
self._connect()
def _connect(self):
if not self.db_url:
logger.error("No NEON_DATABASE_URL provided. Cannot connect to DB.")
return
try:
self.conn = psycopg2.connect(self.db_url)
self._initialize_tables()
logger.info("Connected to PostgreSQL DB.")
except Exception as e:
logger.error(f"Failed to connect to database: {e}.")
self.conn = None
def _ensure_connection(self):
if self.db_url:
if self.conn is None or self.conn.closed != 0:
self._connect()
def _initialize_tables(self):
if not self.conn:
return
try:
with self.conn.cursor() as cur:
cur.execute("""
CREATE TABLE IF NOT EXISTS sessions (
id VARCHAR(255) PRIMARY KEY,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
metadata TEXT
);
""")
cur.execute("""
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS ip_address TEXT;
ALTER TABLE sessions ADD COLUMN IF NOT EXISTS location_data TEXT;
""")
cur.execute("""
CREATE TABLE IF NOT EXISTS conversations (
id VARCHAR(255) PRIMARY KEY,
session_id VARCHAR(255) REFERENCES sessions(id),
user_message TEXT,
ai_response TEXT,
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
""")
self.conn.commit()
except Exception as e:
logger.error(f"Error initializing tables: {e}")
try:
self.conn.rollback()
except Exception:
pass
def create_session(self, metadata="", ip_address=None, location_data=None):
self._ensure_connection()
session_id = str(uuid.uuid4())
location_data_str = json.dumps(location_data) if location_data else None
if self.conn:
try:
with self.conn.cursor() as cur:
cur.execute(
"INSERT INTO sessions (id, metadata, ip_address, location_data) VALUES (%s, %s, %s, %s)",
(session_id, metadata, ip_address, location_data_str)
)
self.conn.commit()
except Exception as e:
logger.error(f"DB Error creating session: {e}")
try:
self.conn.rollback()
except Exception:
self.conn = None
return session_id
def add_message(self, session_id, user_message, ai_response):
self._ensure_connection()
conv_id = str(uuid.uuid4())
if self.conn:
try:
with self.conn.cursor() as cur:
cur.execute(
"INSERT INTO conversations (id, session_id, user_message, ai_response) VALUES (%s, %s, %s, %s)",
(conv_id, session_id, user_message, ai_response)
)
self.conn.commit()
except Exception as e:
logger.error(f"DB Error adding message: {e}")
try:
self.conn.rollback()
except Exception:
self.conn = None
def get_history(self, session_id):
self._ensure_connection()
if self.conn:
try:
with self.conn.cursor() as cur:
cur.execute(
"SELECT user_message, ai_response, timestamp FROM conversations WHERE session_id = %s ORDER BY timestamp ASC",
(session_id,)
)
rows = cur.fetchall()
return [{"user_message": row[0], "ai_response": row[1], "timestamp": row[2].isoformat()} for row in rows]
except Exception as e:
logger.error(f"DB Error getting history: {e}")
try:
self.conn.rollback()
except Exception:
self.conn = None
return []
|