Spaces:
Running on Zero
Running on Zero
Download backend/database.py from Lvvbaaal/legalgpt-api: direct link, hf CLI and curl.
- Browser
- Download file 28.4 kB
-
https://huggingface.co/spaces/Lvvbaaal/legalgpt-api/resolve/main/backend/database.py
- Command line
-
hf download hf://spaces/Lvvbaaal/legalgpt-api/backend/database.py
-
curl -L -o database.py https://huggingface.co/spaces/Lvvbaaal/legalgpt-api/resolve/main/backend/database.py
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() | |