File size: 5,188 Bytes
5eb5327 |
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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import sqlite3
from datetime import datetime
DB_PATH = "chatai.db"
def get_conn():
return sqlite3.connect(DB_PATH, check_same_thread=False)
def init_db():
conn = get_conn()
c = conn.cursor()
c.execute("""
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT UNIQUE NOT NULL,
password_hash BLOB NOT NULL,
role TEXT NOT NULL DEFAULT 'user',
is_active INTEGER NOT NULL DEFAULT 1,
created_at TEXT NOT NULL
)
""")
c.execute("""
CREATE TABLE IF NOT EXISTS conversations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
title TEXT,
created_at TEXT NOT NULL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
""")
c.execute("""
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
conversation_id INTEGER NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
attachments TEXT,
created_at TEXT NOT NULL,
FOREIGN KEY(conversation_id) REFERENCES conversations(id)
)
""")
conn.commit()
return conn
# --- User CRUD ---
def create_user(username, password_hash, role="user", is_active=True):
conn = get_conn()
c = conn.cursor()
c.execute(
"INSERT INTO users (username, password_hash, role, is_active, created_at) VALUES (?,?,?,?,?)",
(username, password_hash, role, 1 if is_active else 0, datetime.utcnow().isoformat()),
)
conn.commit()
return c.lastrowid
def get_user_by_username(username):
conn = get_conn()
c = conn.cursor()
c.execute("SELECT id, username, password_hash, role, is_active, created_at FROM users WHERE username = ?", (username,))
row = c.fetchone()
if not row: return None
return {"id": row[0], "username": row[1], "password_hash": row[2], "role": row[3], "is_active": bool(row[4]), "created_at": row[5]}
def get_user_by_id(user_id):
conn = get_conn()
c = conn.cursor()
c.execute("SELECT id, username, password_hash, role, is_active, created_at FROM users WHERE id = ?", (user_id,))
row = c.fetchone()
if not row: return None
return {"id": row[0], "username": row[1], "password_hash": row[2], "role": row[3], "is_active": bool(row[4]), "created_at": row[5]}
def list_users():
conn = get_conn()
c = conn.cursor()
c.execute("SELECT id, username, role, is_active, created_at FROM users ORDER BY id ASC")
rows = c.fetchall()
return [{"id": r[0], "username": r[1], "role": r[2], "is_active": bool(r[3]), "created_at": r[4]} for r in rows]
def set_user_active(user_id, active: bool):
conn = get_conn()
c = conn.cursor()
c.execute("UPDATE users SET is_active=? WHERE id=?", (1 if active else 0, user_id))
conn.commit()
def set_user_role(user_id, role: str):
conn = get_conn()
c = conn.cursor()
c.execute("UPDATE users SET role=? WHERE id=?", (role, user_id))
conn.commit()
def update_user_password(user_id, new_password_hash):
conn = get_conn()
c = conn.cursor()
c.execute("UPDATE users SET password_hash=? WHERE id=?", (new_password_hash, user_id))
conn.commit()
def delete_user(user_id):
conn = get_conn()
c = conn.cursor()
c.execute("DELETE FROM users WHERE id=?", (user_id,))
conn.commit()
# --- Conversations & Messages ---
def create_conversation(user_id, title="New Chat"):
conn = get_conn()
c = conn.cursor()
c.execute("INSERT INTO conversations (user_id, title, created_at) VALUES (?,?,?)",
(user_id, title, datetime.utcnow().isoformat()))
conn.commit()
return c.lastrowid
def list_conversations(user_id):
conn = get_conn()
c = conn.cursor()
c.execute("SELECT id, title, created_at FROM conversations WHERE user_id=? ORDER BY id DESC", (user_id,))
rows = c.fetchall()
return [{"id": r[0], "title": r[1], "created_at": r[2]} for r in rows]
def rename_conversation(conversation_id, title):
conn = get_conn()
c = conn.cursor()
c.execute("UPDATE conversations SET title=? WHERE id=?", (title, conversation_id))
conn.commit()
def delete_conversation(conversation_id):
conn = get_conn()
c = conn.cursor()
c.execute("DELETE FROM messages WHERE conversation_id=?", (conversation_id,))
c.execute("DELETE FROM conversations WHERE id=?", (conversation_id,))
conn.commit()
def add_message(conversation_id, role, content, attachments=None):
conn = get_conn()
c = conn.cursor()
c.execute("INSERT INTO messages (conversation_id, role, content, attachments, created_at) VALUES (?,?,?,?,?)",
(conversation_id, role, content, attachments or "[]", datetime.utcnow().isoformat()))
conn.commit()
return c.lastrowid
def get_messages(conversation_id):
conn = get_conn()
c = conn.cursor()
c.execute("SELECT role, content, attachments, created_at FROM messages WHERE conversation_id=? ORDER BY id ASC",
(conversation_id,))
rows = c.fetchall()
msgs = []
for r in rows:
msgs.append({"role": r[0], "content": r[1], "attachments": r[2], "created_at": r[3]})
return msgs
|