1344 lines
45 KiB
Python
1344 lines
45 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 sqlite3
|
|
import threading
|
|
from datetime import datetime, timezone
|
|
|
|
logger = logging.getLogger(__name__)
|
|
from typing import Any, Iterable, Optional
|
|
|
|
|
|
from utils.paths import 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
|
|
|
|
|
|
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_threads (
|
|
id TEXT NOT NULL PRIMARY KEY,
|
|
title TEXT NOT NULL,
|
|
model_type TEXT NOT NULL,
|
|
model_id TEXT,
|
|
pair_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
|
|
)
|
|
"""
|
|
)
|
|
chat_thread_cols = {
|
|
row[1] for row in conn.execute("PRAGMA table_info(chat_threads)").fetchall()
|
|
}
|
|
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_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 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)"
|
|
)
|
|
# 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,
|
|
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")
|
|
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_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
|
|
|
|
|
|
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,
|
|
"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_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, 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,
|
|
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"),
|
|
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")),
|
|
"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,
|
|
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 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()
|
|
|
|
|
|
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 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()
|