File size: 4,231 Bytes
18e58fa | 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 | import sqlite3
import json
import os
from datetime import datetime
from sqlite_utils import Database
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "lexrag.db")
def setup_db():
db = Database(DB_PATH)
if "conversations" not in db.table_names():
db["conversations"].create({
"id": int,
"session_id": str,
"role": str,
"content": str,
"sources": str,
"timestamp": str,
"provider": str
}, pk="id")
if "sessions" not in db.table_names():
db["sessions"].create({
"session_id": str,
"name": str,
"created_at": str,
"updated_at": str
}, pk="session_id")
# Add indexes if not exists to optimize session lookup performance
try:
db["conversations"].create_index(["session_id"], if_not_exists=True)
db["conversations"].create_index(["timestamp"], if_not_exists=True)
except Exception:
pass
return db
def update_session_name(session_id: str, name: str):
db = setup_db()
now = datetime.now().isoformat()
try:
exists = list(db["sessions"].rows_where("session_id = ?", [session_id]))
if exists:
db["sessions"].update(session_id, {"name": name, "updated_at": now})
else:
db["sessions"].insert({
"session_id": session_id,
"name": name,
"created_at": now,
"updated_at": now
})
except Exception as e:
print(f"Warning: could not update session name: {e}")
def get_session_name(session_id: str) -> str:
db = setup_db()
try:
rows = list(db["sessions"].rows_where("session_id = ?", [session_id]))
return rows[0]["name"] if rows else session_id[:8]
except Exception:
return session_id[:8]
def save_message(session_id: str, role: str, content: str, sources: list = None, provider: str = None):
db = setup_db()
db["conversations"].insert({
"session_id": session_id,
"role": role,
"content": content,
"sources": json.dumps(sources) if sources else "[]",
"timestamp": datetime.now().isoformat(),
"provider": provider
})
def get_history(session_id: str, limit: int = 10):
db = setup_db()
rows = db["conversations"].rows_where(
"session_id = ?", [session_id], order_by="timestamp DESC", limit=limit
)
history = list(rows)
history.reverse()
return [{"role": r["role"], "content": r["content"]} for r in history]
def get_history_full(session_id: str, limit: int = 100):
"""Returns full message objects including sources for UI rendering."""
db = setup_db()
rows = db["conversations"].rows_where(
"session_id = ?", [session_id], order_by="timestamp ASC", limit=limit
)
result = []
for r in rows:
sources = []
try:
sources = json.loads(r.get("sources", "[]"))
except Exception:
pass
result.append({
"role": r["role"],
"content": r["content"],
"sources": sources,
"provider": r.get("provider"),
"timestamp": r.get("timestamp")
})
return result
def list_sessions():
db = setup_db()
if "conversations" not in db.table_names():
return []
rows = list(db.query("""
SELECT c.session_id,
COALESCE(s.name, c.session_id) as name,
MAX(c.timestamp) as last_active,
MIN(c.timestamp) as created_at,
COUNT(*) as message_count,
(SELECT content FROM conversations c2 WHERE c2.session_id = c.session_id AND c2.role = 'user' ORDER BY c2.timestamp ASC LIMIT 1) as preview
FROM conversations c
LEFT JOIN sessions s ON s.session_id = c.session_id
GROUP BY c.session_id
ORDER BY last_active DESC
"""))
return rows
def delete_session(session_id: str):
db = setup_db()
db["conversations"].delete_where("session_id = ?", [session_id])
try:
db["sessions"].delete(session_id)
except Exception:
pass
|