unsloth/studio/backend/storage/studio_db.py
Roland Tannous c267c3772c Merge remote-tracking branch 'origin/main' into feature/rag
# Conflicts:
#	studio/backend/tests/test_desktop_auth.py
#	studio/frontend/package-lock.json
2026-06-03 10:43:15 +04:00

1715 lines
57 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
SQLite storage for training run history and metrics.
Follows the same pattern as auth/storage.py — module-level functions,
raw sqlite3, per-function connections. Enhancements over auth:
- WAL mode for concurrent read/write access
- PRAGMA foreign_keys = ON for CASCADE deletes
"""
import json
import logging
import os
import platform
import re
import shutil
import sqlite3
import threading
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
logger = logging.getLogger(__name__)
from typing import Any, Iterable, Optional
from utils.paths import project_workspaces_root, studio_db_path, ensure_dir
def _denied_path_prefixes() -> list[str]:
"""Platform-aware denylist of system directories."""
system = platform.system()
if system == "Linux":
return ["/proc", "/sys", "/dev", "/etc", "/boot", "/run"]
if system == "Darwin":
# realpath() resolves /etc -> /private/etc, /tmp -> /private/tmp on macOS,
# so include the /private variants to avoid bypasses.
return [
"/System",
"/Library",
"/dev",
"/etc",
"/private/etc",
"/tmp",
"/private/tmp",
"/var",
"/private/var",
]
if system == "Windows":
win = os.environ.get("SystemRoot", r"C:\Windows")
pf = os.environ.get("ProgramFiles", r"C:\Program Files")
pf86 = os.environ.get("ProgramFiles(x86)", r"C:\Program Files (x86)")
return [os.path.normcase(p) for p in [win, pf, pf86]]
return []
_schema_lock = threading.Lock()
_schema_ready = False
_SQLITE_IN_CHUNK_SIZE = 900
_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",)
def _project_slug(name: str) -> str:
slug = re.sub(r"[^A-Za-z0-9._-]+", "-", name.strip()).strip(".-_")
return slug[:48] or "project"
def _default_project_root(project: dict) -> str:
project_id = str(project["id"])
suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project"
folder_name = f"{_project_slug(str(project.get('name') or 'Project'))}-{suffix}"
return str(project_workspaces_root() / folder_name)
def _ensure_project_workspace(root_path: str) -> str:
root = Path(root_path).expanduser()
root_resolved = ensure_dir(root).resolve()
for subdir in _PROJECT_WORKSPACE_SUBDIRS:
ensure_dir(root_resolved / subdir)
return str(root_resolved)
def _delete_project_workspace(project: dict) -> None:
root_path = project.get("rootPath")
if not root_path:
return
root = Path(root_path).expanduser()
try:
root_resolved = root.resolve(strict = False)
except (OSError, RuntimeError, ValueError):
logger.warning(
"Skipping project workspace delete for invalid path %r", root_path
)
return
project_id = str(project["id"])
suffix = re.sub(r"[^A-Za-z0-9_-]+", "-", project_id)[:8].strip("-_") or "project"
if not root_resolved.name.endswith(f"-{suffix}"):
logger.warning(
"Skipping project workspace delete for unexpected project path %s",
root_resolved,
)
return
if root_resolved.parent == root_resolved or root_resolved == Path.home().resolve():
logger.warning(
"Skipping project workspace delete for unsafe project path %s",
root_resolved,
)
return
check = (
os.path.normcase(str(root_resolved))
if platform.system() == "Windows"
else str(root_resolved)
)
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
logger.warning(
"Skipping project workspace delete under denied path %s",
root_resolved,
)
return
if not root_resolved.exists():
return
if root_resolved.is_symlink() or not root_resolved.is_dir():
logger.warning(
"Skipping project workspace delete for non-directory path %s",
root_resolved,
)
return
shutil.rmtree(root_resolved)
def _ensure_schema(conn: sqlite3.Connection) -> None:
"""Create tables and indexes if they don't exist. Called once per process."""
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_runs (
id TEXT NOT NULL PRIMARY KEY,
status TEXT NOT NULL DEFAULT 'running',
model_name TEXT NOT NULL,
dataset_name TEXT NOT NULL,
config_json TEXT NOT NULL,
started_at TEXT NOT NULL,
ended_at TEXT,
total_steps INTEGER,
final_step INTEGER,
final_loss REAL,
output_dir TEXT,
error_message TEXT,
duration_seconds REAL,
loss_sparkline TEXT,
display_name TEXT
)
"""
)
existing_cols = {
row[1] for row in conn.execute("PRAGMA table_info(training_runs)").fetchall()
}
if "display_name" not in existing_cols:
conn.execute("ALTER TABLE training_runs ADD COLUMN display_name TEXT")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS training_metrics (
id INTEGER PRIMARY KEY AUTOINCREMENT,
run_id TEXT NOT NULL REFERENCES training_runs(id) ON DELETE CASCADE,
step INTEGER NOT NULL,
loss REAL,
learning_rate REAL,
grad_norm REAL,
eval_loss REAL,
epoch REAL,
num_tokens INTEGER,
elapsed_seconds REAL,
UNIQUE(run_id, step)
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_metrics_run_id ON training_metrics(run_id)"
)
# Use COLLATE NOCASE on Windows so C:\Models and c:\models dedup via the
# UNIQUE constraint. On Linux/macOS (case-sensitive FS) keep the default
# BINARY collation so /Models and /models remain distinct.
collation = "COLLATE NOCASE" if platform.system() == "Windows" else ""
conn.execute(
f"""
CREATE TABLE IF NOT EXISTS scan_folders (
id INTEGER PRIMARY KEY AUTOINCREMENT,
path TEXT NOT NULL UNIQUE {collation},
created_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_projects (
id TEXT NOT NULL PRIMARY KEY,
name TEXT NOT NULL,
instructions TEXT,
root_path TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
updated_at INTEGER NOT NULL
)
"""
)
chat_project_cols = {
row[1] for row in conn.execute("PRAGMA table_info(chat_projects)").fetchall()
}
if "root_path" not in chat_project_cols:
conn.execute("ALTER TABLE chat_projects ADD COLUMN root_path TEXT")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_projects_archived_updated_at ON chat_projects(archived, updated_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_threads (
id TEXT NOT NULL PRIMARY KEY,
title TEXT NOT NULL,
model_type TEXT NOT NULL,
model_id TEXT,
pair_id TEXT,
project_id TEXT,
archived INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
openai_code_exec_container_id TEXT,
anthropic_code_exec_container_id TEXT,
FOREIGN KEY(project_id) REFERENCES chat_projects(id) ON DELETE CASCADE
)
"""
)
chat_thread_cols = {
row[1] for row in conn.execute("PRAGMA table_info(chat_threads)").fetchall()
}
if "project_id" not in chat_thread_cols:
conn.execute("ALTER TABLE chat_threads ADD COLUMN project_id TEXT")
if "openai_code_exec_container_id" not in chat_thread_cols:
conn.execute(
"ALTER TABLE chat_threads ADD COLUMN openai_code_exec_container_id TEXT"
)
if "anthropic_code_exec_container_id" not in chat_thread_cols:
conn.execute(
"ALTER TABLE chat_threads ADD COLUMN anthropic_code_exec_container_id TEXT"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_messages (
id TEXT NOT NULL PRIMARY KEY,
thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
parent_id TEXT,
role TEXT NOT NULL,
content_json TEXT NOT NULL,
attachments_json TEXT,
metadata_json TEXT,
created_at INTEGER NOT NULL
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_pair_id ON chat_threads(pair_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_project_id ON chat_threads(project_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_messages_thread_id_created_at ON chat_messages(thread_id, created_at)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_settings (
key TEXT NOT NULL PRIMARY KEY,
value_json TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_settings (
key TEXT NOT NULL PRIMARY KEY,
value_json TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_settings_quarantine (
id INTEGER PRIMARY KEY AUTOINCREMENT,
key TEXT NOT NULL,
value_json TEXT NOT NULL,
reason TEXT NOT NULL,
quarantined_at TEXT NOT NULL
)
"""
)
# Server-side import ledger so a studio.db wipe correctly re-triggers
# the legacy Dexie import. The previous boolean localStorage sentinel
# (`unsloth_chat_legacy_imported_to_studio_db`) is non-recoverable:
# if studio.db is recreated while the browser keeps the flag, legacy
# Dexie threads are silently hidden from the sidebar. The ledger
# lives inside studio.db so it disappears together with the data it
# is supposed to track, which is the recovery the boolean lacked.
# Keyed by legacy thread id; per-thread is sufficient because Dexie
# is read-only after this PR (a thread's message set does not grow).
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_legacy_imports (
legacy_thread_id TEXT NOT NULL PRIMARY KEY,
imported_at INTEGER NOT NULL
) WITHOUT ROWID
"""
)
# RAG schema. rag_documents enforces XOR on (kb_id, thread_id).
# chunking_strategy/mode are immutable post-create (invalidates chunks).
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_knowledge_bases (
id TEXT NOT NULL PRIMARY KEY,
name TEXT NOT NULL UNIQUE,
description TEXT,
owner_user_id TEXT,
embedding_model TEXT NOT NULL,
chunking_strategy TEXT NOT NULL DEFAULT 'standard',
mode TEXT NOT NULL DEFAULT 'text',
created_at INTEGER NOT NULL
)
"""
)
# Idempotent ALTER for pre-existing installs.
kb_cols = {
row[1]
for row in conn.execute("PRAGMA table_info(rag_knowledge_bases)").fetchall()
}
if "chunking_strategy" not in kb_cols:
conn.execute(
"ALTER TABLE rag_knowledge_bases "
"ADD COLUMN chunking_strategy TEXT NOT NULL DEFAULT 'standard'"
)
if "mode" not in kb_cols:
conn.execute(
"ALTER TABLE rag_knowledge_bases "
"ADD COLUMN mode TEXT NOT NULL DEFAULT 'text'"
)
# thread_id has no FK: docs can attach before the thread is persisted.
# chat_history DELETE handlers purge matching rag_documents explicitly.
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_documents (
id TEXT NOT NULL PRIMARY KEY,
kb_id TEXT REFERENCES rag_knowledge_bases(id) ON DELETE CASCADE,
thread_id TEXT,
filename TEXT NOT NULL,
content_type TEXT,
stored_path TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'pending',
num_chunks INTEGER NOT NULL DEFAULT 0,
byte_size INTEGER NOT NULL DEFAULT 0,
error TEXT,
created_at INTEGER NOT NULL,
CHECK ((kb_id IS NOT NULL) <> (thread_id IS NOT NULL))
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_documents_kb_id ON rag_documents(kb_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_id ON rag_documents(thread_id)"
)
# content_hash: sha256 of uploaded bytes; skips re-indexing a file already
# in the same scope (kb_id / thread_id).
rag_documents_columns = {
row[1] for row in conn.execute("PRAGMA table_info(rag_documents)").fetchall()
}
if "content_hash" not in rag_documents_columns:
conn.execute("ALTER TABLE rag_documents ADD COLUMN content_hash TEXT")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_documents_kb_hash "
"ON rag_documents(kb_id, content_hash)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_documents_thread_hash "
"ON rag_documents(thread_id, content_hash)"
)
# kind: text|image|caption. linked_chunk_id pairs image↔caption (both null for text).
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_chunks (
id TEXT NOT NULL PRIMARY KEY,
document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE,
chunk_index INTEGER NOT NULL,
text TEXT NOT NULL,
token_count INTEGER NOT NULL DEFAULT 0,
page_number INTEGER,
kind TEXT NOT NULL DEFAULT 'text',
image_path TEXT,
linked_chunk_id TEXT,
source_page_index INTEGER,
page_char_start INTEGER,
page_char_end INTEGER,
line_start INTEGER,
line_end INTEGER,
pdf_regions_json TEXT,
UNIQUE(document_id, chunk_index)
)
"""
)
chunk_cols = {
row[1] for row in conn.execute("PRAGMA table_info(rag_chunks)").fetchall()
}
if "kind" not in chunk_cols:
conn.execute(
"ALTER TABLE rag_chunks ADD COLUMN kind TEXT NOT NULL DEFAULT 'text'"
)
if "image_path" not in chunk_cols:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN image_path TEXT")
if "linked_chunk_id" not in chunk_cols:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN linked_chunk_id TEXT")
for column in (
"source_page_index",
"page_char_start",
"page_char_end",
"line_start",
"line_end",
):
if column not in chunk_cols:
conn.execute(f"ALTER TABLE rag_chunks ADD COLUMN {column} INTEGER")
if "pdf_regions_json" not in chunk_cols:
conn.execute("ALTER TABLE rag_chunks ADD COLUMN pdf_regions_json TEXT")
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_chunks_document_id ON rag_chunks(document_id)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_document_pages (
document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE,
page_index INTEGER NOT NULL,
page_number INTEGER,
text TEXT NOT NULL,
char_count INTEGER NOT NULL DEFAULT 0,
line_count INTEGER NOT NULL DEFAULT 0,
created_at INTEGER NOT NULL,
PRIMARY KEY(document_id, page_index)
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_document_pages_document_id "
"ON rag_document_pages(document_id)"
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS rag_ingestion_jobs (
id TEXT NOT NULL PRIMARY KEY,
document_id TEXT NOT NULL REFERENCES rag_documents(id) ON DELETE CASCADE,
status TEXT NOT NULL DEFAULT 'pending',
progress REAL NOT NULL DEFAULT 0.0,
stage TEXT,
error TEXT,
started_at INTEGER,
finished_at INTEGER
)
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_rag_jobs_document_id ON rag_ingestion_jobs(document_id)"
)
def get_connection() -> sqlite3.Connection:
"""Open studio.db with WAL mode, create tables once per process, enable foreign keys."""
global _schema_ready
db_path = studio_db_path()
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
# foreign_keys is session-scoped, must be set per connection
conn.execute("PRAGMA foreign_keys=ON")
if not _schema_ready:
with _schema_lock:
if not _schema_ready:
try:
_ensure_schema(conn)
_schema_ready = True
except Exception:
conn.close()
raise
return conn
@contextmanager
def closing_connection():
"""`get_connection()` that also closes the connection on exit.
Commits on success and rolls back on error like sqlite3's own context
manager, then always closes the connection. Use this instead of the bare
`with get_connection() as conn:` (which commits but never closes, leaning on
GC to release the handle) so connections are freed deterministically,
matching the explicit conn.close() convention used elsewhere in this module.
"""
conn = get_connection()
try:
with conn:
yield conn
finally:
conn.close()
def create_run(
id: str,
model_name: str,
dataset_name: str,
config_json: str,
started_at: str,
total_steps: Optional[int],
) -> None:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, dataset_name, config_json, started_at, total_steps)
VALUES (?, ?, ?, ?, ?, ?)
""",
(id, model_name, dataset_name, config_json, started_at, total_steps),
)
conn.commit()
finally:
conn.close()
def update_run_total_steps(id: str, total_steps: int) -> None:
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET total_steps = ? WHERE id = ?",
(total_steps, id),
)
conn.commit()
finally:
conn.close()
def update_run_progress(
id: str, step: int, loss: Optional[float], duration_seconds: Optional[float]
) -> None:
"""Update current progress on a running training run (called on each metric flush)."""
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET final_step = ?, final_loss = ?, duration_seconds = ? WHERE id = ?",
(step, loss, duration_seconds, id),
)
conn.commit()
finally:
conn.close()
def finish_run(
id: str,
status: str,
ended_at: str,
final_step: Optional[int],
final_loss: Optional[float],
duration_seconds: Optional[float],
loss_sparkline: Optional[str] = None,
output_dir: Optional[str] = None,
error_message: Optional[str] = None,
) -> None:
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs
SET status = ?, ended_at = ?, final_step = ?, final_loss = ?,
duration_seconds = ?, loss_sparkline = ?, output_dir = ?,
error_message = ?
WHERE id = ?
""",
(
status,
ended_at,
final_step,
final_loss,
duration_seconds,
loss_sparkline,
output_dir,
error_message,
id,
),
)
conn.commit()
finally:
conn.close()
def insert_metrics_batch(run_id: str, metrics: list[dict]) -> None:
if not metrics:
return
conn = get_connection()
try:
conn.executemany(
"""
INSERT INTO training_metrics
(run_id, step, loss, learning_rate, grad_norm, eval_loss, epoch, num_tokens, elapsed_seconds)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(run_id, step) DO UPDATE SET
loss = COALESCE(excluded.loss, loss),
learning_rate = COALESCE(excluded.learning_rate, learning_rate),
grad_norm = COALESCE(excluded.grad_norm, grad_norm),
eval_loss = COALESCE(excluded.eval_loss, eval_loss),
epoch = COALESCE(excluded.epoch, epoch),
num_tokens = COALESCE(excluded.num_tokens, num_tokens),
elapsed_seconds = COALESCE(excluded.elapsed_seconds, elapsed_seconds)
""",
[
(
run_id,
m.get("step"),
m.get("loss"),
m.get("learning_rate"),
m.get("grad_norm"),
m.get("eval_loss"),
m.get("epoch"),
m.get("num_tokens"),
m.get("elapsed_seconds"),
)
for m in metrics
],
)
conn.commit()
finally:
conn.close()
def update_run_display_name(id: str, display_name: Optional[str]) -> None:
conn = get_connection()
try:
conn.execute(
"UPDATE training_runs SET display_name = ? WHERE id = ?",
(display_name, id),
)
conn.commit()
finally:
conn.close()
def list_runs(limit: int = 50, offset: int = 0) -> dict:
conn = get_connection()
try:
total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0]
rows = conn.execute(
"""
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
r.ended_at, r.total_steps, r.final_step, r.final_loss,
r.output_dir, r.duration_seconds, r.error_message,
r.loss_sparkline, r.display_name,
CASE
WHEN r.status = 'stopped'
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
ORDER BY started_at DESC
LIMIT ? OFFSET ?
""",
(limit, offset),
).fetchall()
runs = []
for row in rows:
run = dict(row)
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug(
"Failed to parse loss_sparkline for run %s", run.get("id")
)
run["loss_sparkline"] = None
runs.append(run)
return {"runs": runs, "total": total}
finally:
conn.close()
def get_run(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT r.*,
CASE
WHEN r.status = 'stopped'
AND r.output_dir IS NOT NULL
AND EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
THEN 1 ELSE 0
END AS resumed_later
FROM training_runs r
WHERE r.id = ?
""",
(id,),
).fetchone()
if row is None:
return None
run = dict(row)
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug("Failed to parse loss_sparkline for run %s", id)
run["loss_sparkline"] = None
return run
finally:
conn.close()
def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT r.*,
0 AS resumed_later
FROM training_runs r
WHERE r.output_dir = ?
AND r.status = 'stopped'
AND NOT EXISTS (
SELECT 1
FROM training_runs newer
WHERE newer.output_dir = r.output_dir
AND newer.status IN ('stopped', 'completed')
AND newer.started_at > r.started_at
)
ORDER BY r.started_at DESC
LIMIT 1
""",
(output_dir,),
).fetchone()
if row is None:
return None
run = dict(row)
sparkline = run.get("loss_sparkline")
if sparkline:
try:
run["loss_sparkline"] = json.loads(sparkline)
except (json.JSONDecodeError, TypeError):
logger.debug(
"Failed to parse loss_sparkline for output_dir %s", output_dir
)
run["loss_sparkline"] = None
return run
finally:
conn.close()
def get_run_metrics(id: str) -> dict:
"""Return metric arrays for a run, using paired step arrays per metric."""
conn = get_connection()
try:
rows = conn.execute(
"""
SELECT step, loss, learning_rate, grad_norm, eval_loss, epoch,
num_tokens, elapsed_seconds
FROM training_metrics
WHERE run_id = ?
ORDER BY step
""",
(id,),
).fetchall()
step_history: list[int] = []
loss_history: list[float] = []
loss_step_history: list[int] = []
lr_history: list[float] = []
lr_step_history: list[int] = []
grad_norm_history: list[float] = []
grad_norm_step_history: list[int] = []
eval_loss_history: list[float] = []
eval_step_history: list[int] = []
final_epoch: float | None = None
final_num_tokens: int | None = None
for row in rows:
step = row["step"]
step_history.append(step)
if step > 0 and row["loss"] is not None:
loss_history.append(row["loss"])
loss_step_history.append(step)
if step > 0 and row["learning_rate"] is not None:
lr_history.append(row["learning_rate"])
lr_step_history.append(step)
if step > 0 and row["grad_norm"] is not None:
grad_norm_history.append(row["grad_norm"])
grad_norm_step_history.append(step)
if step > 0 and row["eval_loss"] is not None:
eval_loss_history.append(row["eval_loss"])
eval_step_history.append(step)
if row["epoch"] is not None:
final_epoch = row["epoch"]
if row["num_tokens"] is not None:
final_num_tokens = row["num_tokens"]
return {
"step_history": step_history,
"loss_history": loss_history,
"loss_step_history": loss_step_history,
"lr_history": lr_history,
"lr_step_history": lr_step_history,
"grad_norm_history": grad_norm_history,
"grad_norm_step_history": grad_norm_step_history,
"eval_loss_history": eval_loss_history,
"eval_step_history": eval_step_history,
"final_epoch": final_epoch,
"final_num_tokens": final_num_tokens,
}
finally:
conn.close()
def delete_run(id: str) -> None:
conn = get_connection()
try:
conn.execute("DELETE FROM training_runs WHERE id = ?", (id,))
conn.commit()
finally:
conn.close()
def cleanup_orphaned_runs() -> None:
"""Mark any 'running' rows as errored on startup (server restarted mid-training)."""
conn = get_connection()
try:
conn.execute(
"""
UPDATE training_runs
SET status = 'error',
error_message = 'Server restarted during training',
ended_at = ?
WHERE status = 'running'
""",
(datetime.now(timezone.utc).isoformat(),),
)
conn.commit()
finally:
conn.close()
def list_scan_folders() -> list[dict]:
conn = get_connection()
try:
rows = conn.execute(
"SELECT id, path, created_at FROM scan_folders ORDER BY created_at"
).fetchall()
return [dict(row) for row in rows]
finally:
conn.close()
def add_scan_folder(path: str) -> dict:
"""Add a directory to the custom scan folder list. Returns the row."""
if not path or not path.strip():
raise ValueError("Path cannot be empty")
normalized = os.path.realpath(os.path.expanduser(path.strip()))
# Validate the path is an existing, readable directory before persisting.
if not os.path.exists(normalized):
raise ValueError("Path does not exist")
if not os.path.isdir(normalized):
raise ValueError("Path must be a directory, not a file")
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
# On Windows, use normcase for denylist comparison but store the
# original-cased path so downstream consumers see the native
# drive-letter casing the user expects (e.g. C:\Models, not c:\models).
is_win = platform.system() == "Windows"
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
raise ValueError(f"Path under {prefix} is not allowed")
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
# On Windows, use case-insensitive lookup so C:\Models and c:\models
# dedup correctly while preserving the originally-stored casing.
if is_win:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE",
(normalized,),
).fetchone()
else:
existing = conn.execute(
"SELECT id, path, created_at FROM scan_folders WHERE path = ?",
(normalized,),
).fetchone()
if existing is not None:
return dict(existing)
try:
conn.execute(
"INSERT INTO scan_folders (path, created_at) VALUES (?, ?)",
(normalized, now),
)
conn.commit()
except sqlite3.IntegrityError:
pass # duplicate -- fall through to SELECT
# Use the same collation as the pre-check so we find the row even
# when a concurrent writer stored it with different casing (Windows).
fallback_sql = (
"SELECT id, path, created_at FROM scan_folders WHERE path = ? COLLATE NOCASE"
if is_win
else "SELECT id, path, created_at FROM scan_folders WHERE path = ?"
)
row = conn.execute(fallback_sql, (normalized,)).fetchone()
if row is None:
raise ValueError("Folder was concurrently removed")
return dict(row)
finally:
conn.close()
def remove_scan_folder(id: int) -> None:
conn = get_connection()
try:
conn.execute("DELETE FROM scan_folders WHERE id = ?", (id,))
conn.commit()
finally:
conn.close()
def _json_loads(value: str | None, fallback):
if value is None:
return fallback
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return fallback
def _chat_thread_from_row(row: sqlite3.Row) -> dict:
data = dict(row)
return {
"id": data["id"],
"title": data["title"],
"modelType": data["model_type"],
"modelId": data.get("model_id") or "",
"pairId": data.get("pair_id") or None,
"projectId": data.get("project_id") or None,
"archived": bool(data["archived"]),
"createdAt": data["created_at"],
"openaiCodeExecContainerId": data.get("openai_code_exec_container_id"),
"anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"),
}
def _chat_project_from_row(row: sqlite3.Row) -> dict:
data = dict(row)
root_path = data.get("root_path")
return {
"id": data["id"],
"name": data["name"],
"instructions": data.get("instructions") or "",
"rootPath": root_path or None,
"sandboxPath": os.path.join(root_path, "sandbox") if root_path else None,
"archived": bool(data["archived"]),
"createdAt": data["created_at"],
"updatedAt": data["updated_at"],
}
def _chat_message_from_row(row: sqlite3.Row) -> dict:
data = dict(row)
message = {
"id": data["id"],
"threadId": data["thread_id"],
"parentId": data.get("parent_id"),
"role": data["role"],
"content": _json_loads(data.get("content_json"), []),
"createdAt": data["created_at"],
}
attachments = _json_loads(data.get("attachments_json"), None)
metadata = _json_loads(data.get("metadata_json"), None)
if attachments is not None:
message["attachments"] = attachments
if metadata is not None:
message["metadata"] = metadata
return message
def upsert_chat_thread(thread: dict) -> dict:
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO chat_threads
(id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
title = excluded.title,
model_type = excluded.model_type,
model_id = excluded.model_id,
pair_id = excluded.pair_id,
project_id = excluded.project_id,
archived = excluded.archived,
created_at = excluded.created_at,
openai_code_exec_container_id = excluded.openai_code_exec_container_id,
anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id
""",
(
thread["id"],
thread.get("title") or "New Chat",
thread["modelType"],
thread.get("modelId") or "",
thread.get("pairId"),
thread.get("projectId"),
1 if thread.get("archived") else 0,
int(thread["createdAt"]),
thread.get("openaiCodeExecContainerId"),
thread.get("anthropicCodeExecContainerId"),
),
)
conn.commit()
return get_chat_thread(thread["id"]) or thread
finally:
conn.close()
def update_chat_thread(id: str, patch: dict) -> Optional[dict]:
allowed = {
"title": ("title", patch.get("title")),
"modelType": ("model_type", patch.get("modelType")),
"modelId": ("model_id", patch.get("modelId")),
"pairId": ("pair_id", patch.get("pairId")),
"projectId": ("project_id", patch.get("projectId")),
"archived": ("archived", 1 if patch.get("archived") else 0),
"createdAt": ("created_at", patch.get("createdAt")),
"openaiCodeExecContainerId": (
"openai_code_exec_container_id",
patch.get("openaiCodeExecContainerId"),
),
"anthropicCodeExecContainerId": (
"anthropic_code_exec_container_id",
patch.get("anthropicCodeExecContainerId"),
),
}
assignments = []
values = []
for key, (column, value) in allowed.items():
if key in patch:
assignments.append(f"{column} = ?")
values.append(value)
if not assignments:
return get_chat_thread(id)
conn = get_connection()
try:
conn.execute(
f"UPDATE chat_threads SET {', '.join(assignments)} WHERE id = ?",
(*values, id),
)
conn.commit()
row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone()
return _chat_thread_from_row(row) if row is not None else None
finally:
conn.close()
def get_chat_thread(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM chat_threads WHERE id = ?", (id,)).fetchone()
return _chat_thread_from_row(row) if row is not None else None
finally:
conn.close()
def list_chat_threads(
model_type: str | None = None,
pair_id: str | None = None,
project_id: str | None = None,
include_archived: bool = True,
) -> list[dict]:
clauses = []
values: list[object] = []
if model_type is not None:
clauses.append("model_type = ?")
values.append(model_type)
if pair_id is not None:
clauses.append("pair_id = ?")
values.append(pair_id)
if project_id is not None:
clauses.append("project_id = ?")
values.append(project_id)
if not include_archived:
clauses.append("archived = 0")
where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
conn = get_connection()
try:
rows = conn.execute(
f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC",
values,
).fetchall()
return [_chat_thread_from_row(row) for row in rows]
finally:
conn.close()
def delete_chat_threads(ids: list[str]) -> None:
if not ids:
return
conn = get_connection()
try:
conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids])
conn.commit()
finally:
conn.close()
def clear_chat_history() -> None:
conn = get_connection()
try:
conn.execute("DELETE FROM chat_threads")
conn.commit()
finally:
conn.close()
def count_chat_threads() -> int:
conn = get_connection()
try:
return int(conn.execute("SELECT COUNT(*) FROM chat_threads").fetchone()[0])
finally:
conn.close()
def upsert_chat_project(project: dict) -> dict:
existing = get_chat_project(project["id"])
root_path = existing.get("rootPath") if existing else None
if not root_path:
root_path = _default_project_root(project)
root_path = _ensure_project_workspace(root_path)
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO chat_projects
(id, name, instructions, root_path, archived, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name = excluded.name,
instructions = excluded.instructions,
root_path = COALESCE(chat_projects.root_path, excluded.root_path),
archived = excluded.archived,
created_at = excluded.created_at,
updated_at = excluded.updated_at
""",
(
project["id"],
project["name"],
project.get("instructions") or "",
root_path,
1 if project.get("archived") else 0,
int(project["createdAt"]),
int(project["updatedAt"]),
),
)
conn.commit()
return get_chat_project(project["id"]) or project
finally:
conn.close()
def update_chat_project(id: str, patch: dict) -> Optional[dict]:
allowed = {
"name": ("name", patch.get("name")),
"instructions": ("instructions", patch.get("instructions")),
"archived": ("archived", 1 if patch.get("archived") else 0),
"createdAt": ("created_at", patch.get("createdAt")),
"updatedAt": ("updated_at", patch.get("updatedAt")),
}
assignments = []
values = []
for key, (column, value) in allowed.items():
if key in patch:
assignments.append(f"{column} = ?")
values.append(value)
if not assignments:
return get_chat_project(id)
conn = get_connection()
try:
conn.execute(
f"UPDATE chat_projects SET {', '.join(assignments)} WHERE id = ?",
(*values, id),
)
conn.commit()
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
return _chat_project_from_row(row) if row is not None else None
finally:
conn.close()
def ensure_chat_project_workspace(id: str) -> Optional[dict]:
project = get_chat_project(id)
if project is None:
return None
root_path = project.get("rootPath") or _default_project_root(project)
root_path = _ensure_project_workspace(root_path)
if project.get("rootPath") == root_path:
return project
conn = get_connection()
try:
conn.execute(
"UPDATE chat_projects SET root_path = ? WHERE id = ?",
(root_path, id),
)
conn.commit()
finally:
conn.close()
return get_chat_project(id)
def get_chat_project(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
return _chat_project_from_row(row) if row is not None else None
finally:
conn.close()
def list_chat_projects(include_archived: bool = False) -> list[dict]:
conn = get_connection()
try:
where = "" if include_archived else "WHERE archived = 0"
rows = conn.execute(
f"SELECT * FROM chat_projects {where} ORDER BY updated_at DESC"
).fetchall()
return [_chat_project_from_row(row) for row in rows]
finally:
conn.close()
def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
if row is None:
conn.rollback()
return None
project = _chat_project_from_row(row)
conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,))
conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,))
conn.commit()
if delete_files:
_delete_project_workspace(project)
return project
except Exception:
conn.rollback()
raise
finally:
conn.close()
class ChatMessageConflictError(RuntimeError):
"""Raised when a chat message id already belongs to another thread."""
class CorruptSettingsError(RuntimeError):
"""Raised when a partial settings patch would overwrite corrupt settings."""
def _parse_chat_setting_json(key: str, value_json: str) -> tuple[bool, Any]:
try:
return True, json.loads(value_json)
except (json.JSONDecodeError, TypeError) as exc:
logger.warning(
"Corrupt chat_settings JSON; quarantining key=%s error=%s",
key,
exc,
)
return False, None
def _load_chat_settings_for_merge(
conn: sqlite3.Connection,
) -> tuple[dict[str, Any], set[str]]:
rows = conn.execute("SELECT key, value_json FROM chat_settings").fetchall()
current: dict[str, Any] = {}
corrupt: set[str] = set()
now = datetime.now(timezone.utc).isoformat()
for row in rows:
ok, value = _parse_chat_setting_json(row["key"], row["value_json"])
if ok:
current[row["key"]] = value
continue
corrupt.add(row["key"])
conn.execute(
"""
INSERT INTO chat_settings_quarantine
(key, value_json, reason, quarantined_at)
VALUES (?, ?, ?, ?)
""",
(row["key"], row["value_json"], "json_decode_error", now),
)
conn.execute(
"DELETE FROM chat_settings WHERE key = ? AND value_json = ?",
(row["key"], row["value_json"]),
)
return current, corrupt
def _raise_if_chat_message_thread_conflicts(
conn: sqlite3.Connection,
thread_id: str,
message_ids: list[str],
) -> None:
unique_ids = list(dict.fromkeys(message_ids))
if not unique_ids:
return
conflicts: list[str] = []
for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT id FROM chat_messages
WHERE id IN ({placeholders}) AND thread_id != ?
ORDER BY id
""",
(*chunk, thread_id),
).fetchall()
conflicts.extend(row["id"] for row in rows)
if conflicts:
preview = ", ".join(conflicts[:5])
suffix = "" if len(conflicts) <= 5 else f" (+{len(conflicts) - 5} more)"
raise ChatMessageConflictError(
f"Message id already belongs to another thread: {preview}{suffix}"
)
def upsert_chat_message(message: dict) -> dict:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_raise_if_chat_message_thread_conflicts(
conn,
message["threadId"],
[message["id"]],
)
conn.execute(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
parent_id = excluded.parent_id,
role = excluded.role,
content_json = excluded.content_json,
attachments_json = excluded.attachments_json,
metadata_json = excluded.metadata_json,
created_at = excluded.created_at
WHERE excluded.thread_id = chat_messages.thread_id
""",
(
message["id"],
message["threadId"],
message.get("parentId"),
message["role"],
json.dumps(message.get("content", [])),
json.dumps(message.get("attachments"))
if message.get("attachments") is not None
else None,
json.dumps(message.get("metadata"))
if message.get("metadata") is not None
else None,
int(message["createdAt"]),
),
)
conn.commit()
return message
except Exception:
conn.rollback()
raise
finally:
conn.close()
def sync_chat_messages(
thread_id: str,
messages: list[dict],
prune_missing: bool = False,
) -> list[dict]:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_raise_if_chat_message_thread_conflicts(
conn,
thread_id,
[m["id"] for m in messages],
)
if prune_missing:
conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,))
conn.executemany(
"""
INSERT INTO chat_messages
(id, thread_id, parent_id, role, content_json, attachments_json, metadata_json, created_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
parent_id = excluded.parent_id,
role = excluded.role,
content_json = excluded.content_json,
attachments_json = excluded.attachments_json,
metadata_json = excluded.metadata_json,
created_at = excluded.created_at
WHERE excluded.thread_id = chat_messages.thread_id
""",
[
(
m["id"],
thread_id,
m.get("parentId"),
m["role"],
json.dumps(m.get("content", [])),
json.dumps(m.get("attachments"))
if m.get("attachments") is not None
else None,
json.dumps(m.get("metadata"))
if m.get("metadata") is not None
else None,
int(m["createdAt"]),
)
for m in messages
],
)
conn.commit()
return list_chat_messages(thread_id)
except ChatMessageConflictError:
conn.rollback()
raise
except sqlite3.Error:
logger.exception("Failed to sync chat messages for thread %s", thread_id)
conn.rollback()
raise
finally:
conn.close()
def list_chat_messages(thread_id: str) -> list[dict]:
conn = get_connection()
try:
rows = conn.execute(
"""
SELECT * FROM chat_messages
WHERE thread_id = ?
ORDER BY created_at ASC, id ASC
""",
(thread_id,),
).fetchall()
return [_chat_message_from_row(row) for row in rows]
finally:
conn.close()
def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute(
"""
SELECT * FROM chat_messages
WHERE thread_id = ? AND id = ?
""",
(thread_id, message_id),
).fetchone()
return _chat_message_from_row(row) if row is not None else None
finally:
conn.close()
def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]:
if not thread_ids:
return []
unique_thread_ids = list(dict.fromkeys(thread_ids))
messages: list[dict] = []
conn = get_connection()
try:
for start in range(0, len(unique_thread_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_thread_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT * FROM chat_messages
WHERE thread_id IN ({placeholders})
ORDER BY created_at ASC, id ASC
""",
chunk,
).fetchall()
messages.extend(_chat_message_from_row(row) for row in rows)
return sorted(
messages,
key = lambda message: (message["createdAt"], message["id"]),
)
finally:
conn.close()
def get_app_setting(key: str, fallback = None):
conn = get_connection()
try:
row = conn.execute(
"SELECT value_json FROM app_settings WHERE key = ?", (key,)
).fetchone()
if row is None:
return fallback
return _json_loads(row["value_json"], fallback)
finally:
conn.close()
def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
if not settings:
return {}
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO app_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in settings.items()],
)
conn.commit()
rows = conn.execute(
"SELECT key, value_json FROM app_settings ORDER BY key"
).fetchall()
return {row["key"]: _json_loads(row["value_json"], None) for row in rows}
finally:
conn.close()
def list_chat_settings() -> dict[str, Any]:
conn = get_connection()
try:
rows = conn.execute(
"SELECT key, value_json FROM chat_settings ORDER BY key"
).fetchall()
settings: dict[str, Any] = {}
for row in rows:
settings[row["key"]] = _json_loads(row["value_json"], None)
return settings
finally:
conn.close()
def upsert_chat_settings(settings: dict[str, Any]) -> dict[str, Any]:
if not settings:
return list_chat_settings()
conn = get_connection()
try:
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO chat_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in settings.items()],
)
conn.commit()
return list_chat_settings()
finally:
conn.close()
def _deep_merge_settings(
current: dict[str, Any], updates: dict[str, Any]
) -> dict[str, Any]:
merged = dict(current)
for key, value in updates.items():
current_value = merged.get(key)
if isinstance(current_value, dict) and isinstance(value, dict):
merged[key] = _deep_merge_settings(current_value, value)
else:
merged[key] = value
return merged
def upsert_chat_settings_merge(updates: dict[str, Any]) -> dict[str, Any]:
"""Atomic read-merge-write under BEGIN IMMEDIATE so two concurrent writers
cannot drop one another's updates."""
if not updates:
return list_chat_settings()
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
current, corrupt = _load_chat_settings_for_merge(conn)
unsafe_partial_keys = [
key
for key, value in updates.items()
if key in corrupt and isinstance(value, dict)
]
if unsafe_partial_keys:
conn.commit()
keys = ", ".join(sorted(unsafe_partial_keys))
raise CorruptSettingsError(
f"Cannot apply partial settings patch to corrupt key(s): {keys}"
)
merged = _deep_merge_settings(current, updates)
now = datetime.now(timezone.utc).isoformat()
conn.executemany(
"""
INSERT INTO chat_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
[(key, json.dumps(value), now) for key, value in merged.items()],
)
conn.commit()
return merged
except CorruptSettingsError:
raise
except Exception:
conn.rollback()
raise
finally:
conn.close()
# ---------------------------------------------------------------------------
# Legacy Dexie import ledger
# ---------------------------------------------------------------------------
# See the schema comment in _ensure_schema() for the recovery rationale.
def list_chat_legacy_imports() -> list[str]:
"""Return the legacy_thread_id of every thread already imported.
Cheap: scans a single small PK-only table. The frontend stuffs the
result into a Set before walking Dexie, so the diff is O(|Dexie|).
"""
conn = get_connection()
try:
rows = conn.execute(
"SELECT legacy_thread_id FROM chat_legacy_imports"
).fetchall()
return [row[0] for row in rows]
finally:
conn.close()
def upsert_chat_legacy_imports(legacy_thread_ids: list[str]) -> tuple[int, int]:
"""Mark each given legacy thread id as imported. Idempotent.
Returns (accepted, inserted):
- accepted: number of non-empty deduped input ids
- inserted: number of rows that were actually new (not already in ledger)
ON CONFLICT DO NOTHING keeps the existing imported_at when an id is
recorded twice. INSERT...RETURNING reports only the rows that were
actually inserted, so callers can distinguish first-time imports
from idempotent re-runs without an extra SELECT.
"""
ids = list(dict.fromkeys(tid for tid in legacy_thread_ids if tid))
if not ids:
return 0, 0
ts = int(datetime.now(timezone.utc).timestamp() * 1000)
conn = get_connection()
try:
inserted = 0
for tid in ids:
row = conn.execute(
"""
INSERT INTO chat_legacy_imports (legacy_thread_id, imported_at)
VALUES (?, ?)
ON CONFLICT(legacy_thread_id) DO NOTHING
RETURNING legacy_thread_id
""",
(tid, ts),
).fetchone()
if row is not None:
inserted += 1
conn.commit()
return len(ids), inserted
finally:
conn.close()