114 lines
3.4 KiB
JavaScript
114 lines
3.4 KiB
JavaScript
import fs from 'node:fs';
|
|
import path from 'node:path';
|
|
import Database from 'better-sqlite3';
|
|
|
|
export function openDb(sqlitePath) {
|
|
const dir = path.dirname(sqlitePath);
|
|
fs.mkdirSync(dir, { recursive: true });
|
|
const db = new Database(sqlitePath);
|
|
db.pragma('journal_mode = WAL');
|
|
db.pragma('synchronous = NORMAL');
|
|
db.pragma('foreign_keys = ON');
|
|
migrate(db);
|
|
return db;
|
|
}
|
|
|
|
function migrate(db) {
|
|
db.exec(`
|
|
CREATE TABLE IF NOT EXISTS chat_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
request_id TEXT NOT NULL UNIQUE,
|
|
chat_id TEXT,
|
|
ts_request INTEGER NOT NULL,
|
|
ts_first_token INTEGER,
|
|
ts_done INTEGER,
|
|
model TEXT NOT NULL,
|
|
messages_json TEXT NOT NULL,
|
|
user_text TEXT,
|
|
assistant_text TEXT,
|
|
prompt_tokens INTEGER,
|
|
completion_tokens INTEGER,
|
|
total_tokens INTEGER,
|
|
status TEXT NOT NULL,
|
|
error TEXT,
|
|
ip TEXT,
|
|
user_agent TEXT
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_chat_log_ts_request ON chat_log(ts_request DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_chat_log_chat_ts ON chat_log(chat_id, ts_request DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_chat_log_model_ts ON chat_log(model, ts_request DESC);
|
|
CREATE INDEX IF NOT EXISTS idx_chat_log_status_ts ON chat_log(status, ts_request DESC);
|
|
`);
|
|
|
|
// FTS5 (prompt+answer search)
|
|
// content_rowid = chat_log.id
|
|
db.exec(`
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS chat_log_fts USING fts5(
|
|
user_text,
|
|
assistant_text,
|
|
content='chat_log',
|
|
content_rowid='id'
|
|
);
|
|
`);
|
|
|
|
db.exec(`
|
|
CREATE TRIGGER IF NOT EXISTS chat_log_ai AFTER INSERT ON chat_log BEGIN
|
|
INSERT INTO chat_log_fts(rowid, user_text, assistant_text)
|
|
VALUES (new.id, coalesce(new.user_text,''), coalesce(new.assistant_text,''));
|
|
END;
|
|
CREATE TRIGGER IF NOT EXISTS chat_log_ad AFTER DELETE ON chat_log BEGIN
|
|
INSERT INTO chat_log_fts(chat_log_fts, rowid, user_text, assistant_text)
|
|
VALUES('delete', old.id, old.user_text, old.assistant_text);
|
|
END;
|
|
CREATE TRIGGER IF NOT EXISTS chat_log_au AFTER UPDATE ON chat_log BEGIN
|
|
INSERT INTO chat_log_fts(chat_log_fts, rowid, user_text, assistant_text)
|
|
VALUES('delete', old.id, old.user_text, old.assistant_text);
|
|
INSERT INTO chat_log_fts(rowid, user_text, assistant_text)
|
|
VALUES (new.id, coalesce(new.user_text,''), coalesce(new.assistant_text,''));
|
|
END;
|
|
`);
|
|
}
|
|
|
|
export function prepareQueries(db) {
|
|
const insertStart = db.prepare(`
|
|
INSERT INTO chat_log (
|
|
request_id, chat_id, ts_request, model, messages_json, user_text,
|
|
status, ip, user_agent
|
|
) VALUES (
|
|
@request_id, @chat_id, @ts_request, @model, @messages_json, @user_text,
|
|
@status, @ip, @user_agent
|
|
)
|
|
`);
|
|
|
|
const markFirstToken = db.prepare(`
|
|
UPDATE chat_log
|
|
SET ts_first_token = coalesce(ts_first_token, @ts_first_token)
|
|
WHERE request_id = @request_id
|
|
`);
|
|
|
|
const finish = db.prepare(`
|
|
UPDATE chat_log
|
|
SET
|
|
ts_done = @ts_done,
|
|
assistant_text = @assistant_text,
|
|
prompt_tokens = @prompt_tokens,
|
|
completion_tokens = @completion_tokens,
|
|
total_tokens = @total_tokens,
|
|
status = @status,
|
|
error = @error
|
|
WHERE request_id = @request_id
|
|
`);
|
|
|
|
const getByRequestId = db.prepare(`
|
|
SELECT * FROM chat_log WHERE request_id = ?
|
|
`);
|
|
|
|
return {
|
|
insertStart,
|
|
markFirstToken,
|
|
finish,
|
|
getByRequestId,
|
|
};
|
|
}
|