legalgpt-api / backend /database.py
Lvvbaaal's picture
Deploy SaulGPT backend with full ChromaDB, Supabase, and Gradio API server
62be02c verified
Raw History Blame Contribute Delete
28.4 kB
"""
SaulGPT — Database Layer (Hybrid PostgreSQL / SQLite)
======================================================
Unified persistence for users, conversations, and messages.
Automatically uses Supabase PostgreSQL when DATABASE_URL is present,
and cleanly falls back to local SQLite (saulgpt.db) when offline.
Includes connection handling, automatic retry, and FOREIGN KEY safety guards.
"""
import json
import sqlite3
import os
from datetime import datetime
from typing import Optional, List, Dict, Any
from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".env"))
load_dotenv(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env"))
DATABASE_URL = os.getenv("DATABASE_URL", "").strip()
try:
import psycopg2
from psycopg2.extras import RealDictCursor, Json
PSYCOPG2_AVAILABLE = True
except ImportError:
PSYCOPG2_AVAILABLE = False
def is_postgres() -> bool:
return bool(PSYCOPG2_AVAILABLE and DATABASE_URL and DATABASE_URL.startswith("postgres"))
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "saulgpt.db")
def _get_sqlite_conn() -> sqlite3.Connection:
"""Get a SQLite connection with WAL journal mode, busy timeout, and foreign keys."""
conn = sqlite3.connect(DB_PATH, timeout=30.0, check_same_thread=False)
conn.row_factory = sqlite3.Row
try:
conn.execute("PRAGMA busy_timeout=30000")
conn.execute("PRAGMA foreign_keys=ON")
except Exception:
pass
return conn
def _get_pg_conn():
"""Get a PostgreSQL connection to Supabase."""
return psycopg2.connect(DATABASE_URL, cursor_factory=RealDictCursor)
def init_db():
"""Create tables if they don't exist and ensure default guest user."""
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE IF NOT EXISTS users (
id SERIAL PRIMARY KEY,
email TEXT UNIQUE NOT NULL,
username TEXT UNIQUE NOT NULL,
password TEXT NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
);
CREATE TABLE IF NOT EXISTS conversations (
id SERIAL PRIMARY KEY,
user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
title TEXT NOT NULL DEFAULT 'New Chat',
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
);
CREATE TABLE IF NOT EXISTS messages (
id SERIAL PRIMARY KEY,
conv_id INTEGER REFERENCES conversations(id) ON DELETE CASCADE,
conversation_id INTEGER REFERENCES conversations(id) ON DELETE CASCADE,
role TEXT NOT NULL,
content TEXT NOT NULL,
meta JSONB,
turn INTEGER NOT NULL DEFAULT 1,
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_conversations_user ON conversations(user_id, updated_at DESC);
CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conv_id, turn);
""")
cur.execute("""
INSERT INTO users (id, email, username, password)
VALUES (1, 'guest@saulgpt.local', 'Guest User', 'guest_mode_unrestricted')
ON CONFLICT (id) DO NOTHING
""")
conn.commit()
conn.close()
print("[Database] Supabase PostgreSQL connected and initialized.")
return
except Exception as e:
print(f"[Database Warning] Supabase init failed: {e}. Falling back to SQLite.")
# SQLite fallback
conn = _get_sqlite_conn()
try:
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA busy_timeout=30000")
conn.execute("PRAGMA foreign_keys=ON")
conn.executescript("""
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
email TEXT UNIQUE NOT NULL,
username TEXT UNIQUE NOT NULL,
password TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS conversations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
title TEXT NOT NULL DEFAULT 'New Chat',
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
conversation_id INTEGER NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
role TEXT NOT NULL CHECK(role IN ('user','assistant')),
content TEXT NOT NULL,
meta TEXT,
turn INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE INDEX IF NOT EXISTS idx_conversations_user
ON conversations(user_id, updated_at DESC);
CREATE INDEX IF NOT EXISTS idx_messages_conv
ON messages(conversation_id, turn);
""")
conn.execute("""
INSERT OR IGNORE INTO users (id, email, username, password)
VALUES (1, 'guest@saulgpt.local', 'Guest User', 'guest_mode_unrestricted')
""")
conn.commit()
print("[Database] SQLite saulgpt.db initialized.")
except Exception as e:
print(f"[Database] SQLite init_db warning: {e}")
finally:
conn.close()
def ensure_user_exists(user_id: Optional[int] = None) -> int:
"""Ensures a user record exists so FOREIGN KEY constraints never fail."""
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute("""
INSERT INTO users (id, email, username, password)
VALUES (1, 'guest@saulgpt.local', 'Guest User', 'guest_mode_unrestricted')
ON CONFLICT (id) DO NOTHING
""")
conn.commit()
if user_id and isinstance(user_id, int) and user_id > 0:
cur.execute("SELECT id FROM users WHERE id = %s", (user_id,))
row = cur.fetchone()
if row:
conn.close()
return user_id
cur.execute("""
INSERT INTO users (id, email, username, password)
VALUES (%s, %s, %s, %s)
ON CONFLICT (id) DO NOTHING
""", (user_id, f"user_{user_id}@saulgpt.local", f"User {user_id}", "token_authenticated"))
conn.commit()
conn.close()
return user_id
conn.close()
return 1
except Exception as e:
print(f"[Database Postgres] ensure_user_exists warning: {e}")
conn = _get_sqlite_conn()
try:
conn.execute("""
INSERT OR IGNORE INTO users (id, email, username, password)
VALUES (1, 'guest@saulgpt.local', 'Guest User', 'guest_mode_unrestricted')
""")
conn.commit()
if user_id and isinstance(user_id, int) and user_id > 0:
row = conn.execute("SELECT id FROM users WHERE id = ?", (user_id,)).fetchone()
if row:
return user_id
try:
conn.execute(
"INSERT INTO users (id, email, username, password) VALUES (?, ?, ?, ?)",
(user_id, f"user_{user_id}@saulgpt.local", f"User {user_id}", "token_authenticated")
)
conn.commit()
return user_id
except Exception:
pass
return 1
except Exception as e:
print(f"[Database SQLite] ensure_user_exists warning: {e}")
return 1
finally:
conn.close()
def ensure_conversation_exists(conv_id: Optional[int], user_id: int = 1) -> int:
"""Ensures a conversation exists. Guarantees messages never hit FOREIGN KEY errors."""
valid_uid = ensure_user_exists(user_id)
if conv_id and isinstance(conv_id, int) and conv_id > 0:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute("SELECT id FROM conversations WHERE id = %s", (conv_id,))
row = cur.fetchone()
if row:
conn.close()
return conv_id
conn.close()
except Exception as e:
print(f"[Database Postgres] ensure_conversation_exists check error: {e}")
else:
conn = _get_sqlite_conn()
try:
row = conn.execute("SELECT id FROM conversations WHERE id = ?", (conv_id,)).fetchone()
if row:
return conv_id
except Exception as e:
print(f"[Database SQLite] ensure_conversation_exists check error: {e}")
finally:
conn.close()
return create_conversation(valid_uid, "New Chat")
# ─────────────────────────────────────────────────────────────
# USER OPERATIONS
# ─────────────────────────────────────────────────────────────
def create_user(email: str, username: str, password_hash: str) -> Optional[int]:
base_username = (username or email.split("@")[0]).strip()
candidate = base_username
counter = 1
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
while True:
cur.execute("SELECT id FROM users WHERE username = %s AND email != %s", (candidate, email))
if not cur.fetchone():
break
candidate = f"{base_username}_{counter}"
counter += 1
cur.execute(
"INSERT INTO users (email, username, password) VALUES (%s, %s, %s) RETURNING id",
(email.strip().lower(), candidate, password_hash)
)
new_id = cur.fetchone()["id"]
conn.commit()
conn.close()
return new_id
except Exception as e:
print(f"[Database Postgres] create_user error: {e}")
return None
conn = _get_sqlite_conn()
try:
while True:
existing = conn.execute("SELECT id FROM users WHERE username = ? AND email != ?", (candidate, email)).fetchone()
if not existing:
break
candidate = f"{base_username}_{counter}"
counter += 1
cur = conn.execute(
"INSERT INTO users (email, username, password) VALUES (?, ?, ?)",
(email.strip().lower(), candidate, password_hash)
)
conn.commit()
return cur.lastrowid
except Exception as e:
print(f"[Database SQLite] create_user error: {e}")
return None
finally:
conn.close()
def get_user_by_email(email: str) -> Optional[Dict]:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute("SELECT * FROM users WHERE email = %s", (email.strip().lower(),))
row = cur.fetchone()
conn.close()
return dict(row) if row else None
except Exception as e:
print(f"[Database Postgres] get_user_by_email error: {e}")
return None
conn = _get_sqlite_conn()
try:
row = conn.execute("SELECT * FROM users WHERE email = ?", (email.strip().lower(),)).fetchone()
return dict(row) if row else None
except Exception as e:
print(f"[Database SQLite] get_user_by_email error: {e}")
return None
finally:
conn.close()
def get_user_by_id(user_id: int) -> Optional[Dict]:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute("SELECT * FROM users WHERE id = %s", (user_id,))
row = cur.fetchone()
conn.close()
return dict(row) if row else None
except Exception as e:
print(f"[Database Postgres] get_user_by_id error: {e}")
return None
conn = _get_sqlite_conn()
try:
row = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
return dict(row) if row else None
except Exception as e:
print(f"[Database SQLite] get_user_by_id error: {e}")
return None
finally:
conn.close()
# ─────────────────────────────────────────────────────────────
# CONVERSATION OPERATIONS
# ─────────────────────────────────────────────────────────────
def create_conversation(user_id: int, title: str = "New Chat") -> int:
valid_uid = ensure_user_exists(user_id)
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"INSERT INTO conversations (user_id, title) VALUES (%s, %s) RETURNING id",
(valid_uid, title)
)
conv_id = cur.fetchone()["id"]
conn.commit()
conn.close()
return conv_id
except Exception as e:
print(f"[Database Postgres] create_conversation error: {e}")
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"INSERT INTO conversations (user_id, title) VALUES (1, %s) RETURNING id",
(title,)
)
conv_id = cur.fetchone()["id"]
conn.commit()
conn.close()
return conv_id
except Exception as e2:
print(f"[Database Postgres] fallback create_conversation error: {e2}")
return 1
conn = _get_sqlite_conn()
try:
cur = conn.execute(
"INSERT INTO conversations (user_id, title) VALUES (?, ?)",
(valid_uid, title)
)
conn.commit()
return cur.lastrowid
except Exception as e:
print(f"[Database SQLite] create_conversation error: {e}")
try:
cur = conn.execute(
"INSERT INTO conversations (user_id, title) VALUES (1, ?)",
(title,)
)
conn.commit()
return cur.lastrowid
except Exception as e2:
print(f"[Database SQLite] create_conversation fallback error: {e2}")
return 1
finally:
conn.close()
def get_conversations(user_id: int, limit: int = 50, offset: int = 0) -> Dict:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"SELECT id, title, created_at, updated_at FROM conversations "
"WHERE user_id = %s ORDER BY updated_at DESC LIMIT %s OFFSET %s",
(user_id, limit, offset)
)
rows = cur.fetchall()
cur.execute("SELECT COUNT(*) as cnt FROM conversations WHERE user_id = %s", (user_id,))
cnt_row = cur.fetchone()
total = cnt_row["cnt"] if cnt_row else 0
conn.close()
return {"conversations": [dict(r) for r in rows], "total": total}
except Exception as e:
print(f"[Database Postgres] get_conversations error: {e}")
return {"conversations": [], "total": 0}
conn = _get_sqlite_conn()
try:
rows = conn.execute(
"SELECT id, title, created_at, updated_at FROM conversations "
"WHERE user_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?",
(user_id, limit, offset)
).fetchall()
total_row = conn.execute(
"SELECT COUNT(*) as cnt FROM conversations WHERE user_id = ?",
(user_id,)
).fetchone()
total = total_row["cnt"] if total_row else 0
return {"conversations": [dict(r) for r in rows], "total": total}
except Exception as e:
print(f"[Database SQLite] get_conversations error: {e}")
return {"conversations": [], "total": 0}
finally:
conn.close()
def get_conversation(conv_id: int, user_id: Optional[int] = None) -> Optional[Dict]:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
if user_id:
cur.execute(
"SELECT id, title, created_at, updated_at FROM conversations "
"WHERE id = %s AND user_id = %s",
(conv_id, user_id)
)
else:
cur.execute(
"SELECT id, title, created_at, updated_at FROM conversations WHERE id = %s",
(conv_id,)
)
row = cur.fetchone()
conn.close()
return dict(row) if row else None
except Exception as e:
print(f"[Database Postgres] get_conversation error: {e}")
return None
conn = _get_sqlite_conn()
try:
if user_id:
row = conn.execute(
"SELECT id, title, created_at, updated_at FROM conversations "
"WHERE id = ? AND user_id = ?",
(conv_id, user_id)
).fetchone()
else:
row = conn.execute(
"SELECT id, title, created_at, updated_at FROM conversations WHERE id = ?",
(conv_id,)
).fetchone()
return dict(row) if row else None
except Exception as e:
print(f"[Database SQLite] get_conversation error: {e}")
return None
finally:
conn.close()
def update_conversation_title(conv_id: int, title: str):
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"UPDATE conversations SET title = %s, updated_at = NOW() WHERE id = %s",
(title, conv_id)
)
conn.commit()
conn.close()
return
except Exception as e:
print(f"[Database Postgres] update_conversation_title error: {e}")
return
conn = _get_sqlite_conn()
try:
conn.execute(
"UPDATE conversations SET title = ?, updated_at = datetime('now') WHERE id = ?",
(title, conv_id)
)
conn.commit()
except Exception as e:
print(f"[Database SQLite] update_conversation_title error: {e}")
finally:
conn.close()
def touch_conversation(conv_id: int):
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"UPDATE conversations SET updated_at = NOW() WHERE id = %s",
(conv_id,)
)
conn.commit()
conn.close()
return
except Exception as e:
print(f"[Database Postgres] touch_conversation error: {e}")
return
conn = _get_sqlite_conn()
try:
conn.execute(
"UPDATE conversations SET updated_at = datetime('now') WHERE id = ?",
(conv_id,)
)
conn.commit()
except Exception as e:
print(f"[Database SQLite] touch_conversation error: {e}")
finally:
conn.close()
def delete_conversation(conv_id: int, user_id: int) -> bool:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"DELETE FROM conversations WHERE id = %s AND user_id = %s",
(conv_id, user_id)
)
deleted = cur.rowcount > 0
conn.commit()
conn.close()
return deleted
except Exception as e:
print(f"[Database Postgres] delete_conversation error: {e}")
return False
conn = _get_sqlite_conn()
try:
cur = conn.execute(
"DELETE FROM conversations WHERE id = ? AND user_id = ?",
(conv_id, user_id)
)
conn.commit()
return cur.rowcount > 0
except Exception as e:
print(f"[Database SQLite] delete_conversation error: {e}")
return False
finally:
conn.close()
# ─────────────────────────────────────────────────────────────
# MESSAGE OPERATIONS
# ─────────────────────────────────────────────────────────────
def add_message(conv_id: int, role: str, content: str, meta: Any = None, turn: int = 0) -> Optional[int]:
safe_conv_id = ensure_conversation_exists(conv_id)
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
pg_meta = None
if meta is not None:
if isinstance(meta, str):
try:
pg_meta = Json(json.loads(meta))
except Exception:
pg_meta = Json({"raw": meta})
else:
pg_meta = Json(meta)
cur.execute(
"INSERT INTO messages (conv_id, conversation_id, role, content, meta, turn) "
"VALUES (%s, %s, %s, %s, %s, %s) RETURNING id",
(safe_conv_id, safe_conv_id, role, content, pg_meta, turn)
)
msg_id = cur.fetchone()["id"]
conn.commit()
conn.close()
return msg_id
except Exception as e:
print(f"[Database Postgres] add_message error: {e}")
return None
conn = _get_sqlite_conn()
try:
sql_meta = meta
if sql_meta is not None and not isinstance(sql_meta, str):
sql_meta = json.dumps(sql_meta)
cur = conn.execute(
"INSERT INTO messages (conversation_id, role, content, meta, turn) VALUES (?, ?, ?, ?, ?)",
(safe_conv_id, role, content, sql_meta, turn)
)
conn.commit()
return cur.lastrowid
except Exception as e:
print(f"[Database SQLite] add_message error: {e}")
return None
finally:
conn.close()
def get_messages(conv_id: int) -> List[Dict]:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"SELECT role, content, meta, turn, created_at FROM messages "
"WHERE (conv_id = %s OR conversation_id = %s) ORDER BY turn ASC, id ASC",
(conv_id, conv_id)
)
rows = cur.fetchall()
conn.close()
return [dict(r) for r in rows]
except Exception as e:
print(f"[Database Postgres] get_messages error: {e}")
return []
conn = _get_sqlite_conn()
try:
rows = conn.execute(
"SELECT role, content, meta, turn, created_at FROM messages "
"WHERE conversation_id = ? ORDER BY turn ASC",
(conv_id,)
).fetchall()
return [dict(r) for r in rows]
except Exception as e:
print(f"[Database SQLite] get_messages error: {e}")
return []
finally:
conn.close()
def get_last_turn(conv_id: int) -> int:
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
cur.execute(
"SELECT COALESCE(MAX(turn), 0) as last_turn FROM messages "
"WHERE (conv_id = %s OR conversation_id = %s)",
(conv_id, conv_id)
)
row = cur.fetchone()
conn.close()
return row["last_turn"] if row else 0
except Exception as e:
print(f"[Database Postgres] get_last_turn error: {e}")
return 0
conn = _get_sqlite_conn()
try:
row = conn.execute(
"SELECT COALESCE(MAX(turn), 0) as last_turn FROM messages WHERE conversation_id = ?",
(conv_id,)
).fetchone()
return row["last_turn"] if row else 0
except Exception as e:
print(f"[Database SQLite] get_last_turn error: {e}")
return 0
finally:
conn.close()
def bulk_import_messages(conv_id: int, turns: list) -> bool:
"""Import an array of turn dicts into a conversation in one transaction."""
safe_conv_id = ensure_conversation_exists(conv_id)
if is_postgres():
try:
conn = _get_pg_conn()
with conn.cursor() as cur:
for i, t in enumerate(turns):
meta = t.get("meta")
pg_meta = None
if meta is not None:
if isinstance(meta, str):
try:
pg_meta = Json(json.loads(meta))
except Exception:
pg_meta = Json({"raw": meta})
else:
pg_meta = Json(meta)
cur.execute(
"INSERT INTO messages (conv_id, conversation_id, role, content, meta, turn) "
"VALUES (%s, %s, %s, %s, %s, %s)",
(safe_conv_id, safe_conv_id, t["role"], t["content"], pg_meta, i + 1)
)
cur.execute("UPDATE conversations SET updated_at = NOW() WHERE id = %s", (safe_conv_id,))
conn.commit()
conn.close()
return True
except Exception as e:
print(f"[Database Postgres] bulk_import_messages error: {e}")
return False
conn = _get_sqlite_conn()
try:
rows = []
for i, t in enumerate(turns):
meta = t.get("meta")
if meta is not None and not isinstance(meta, str):
meta = json.dumps(meta)
rows.append((safe_conv_id, t["role"], t["content"], meta, i + 1))
conn.executemany(
"INSERT INTO messages (conversation_id, role, content, meta, turn) VALUES (?, ?, ?, ?, ?)",
rows
)
conn.execute(
"UPDATE conversations SET updated_at = datetime('now') WHERE id = ?",
(safe_conv_id,)
)
conn.commit()
return True
except Exception as e:
print(f"[Database SQLite] bulk_import_messages error: {e}")
conn.rollback()
return False
finally:
conn.close()