520 lines
16 KiB
C++
520 lines
16 KiB
C++
#include "chatstore.hpp"
|
|
|
|
#include "llmclient.hpp"
|
|
#include "message.hpp"
|
|
#include "segment.hpp"
|
|
|
|
#include <QDateTime>
|
|
#include <QDebug>
|
|
#include <QDir>
|
|
#include <QFileInfo>
|
|
#include <QStandardPaths>
|
|
#include <QSqlDatabase>
|
|
#include <QSqlError>
|
|
#include <QSqlQuery>
|
|
#include <QUuid>
|
|
|
|
#include <algorithm>
|
|
|
|
namespace ZShell::llm {
|
|
|
|
namespace {
|
|
|
|
QString segmentTypeName(LlmSegment::Type type) {
|
|
switch (type) {
|
|
case LlmSegment::Type::Reasoning:
|
|
return QStringLiteral("reasoning");
|
|
case LlmSegment::Type::ToolCall:
|
|
return QStringLiteral("tool_call");
|
|
case LlmSegment::Type::Content:
|
|
return QStringLiteral("content");
|
|
}
|
|
return QStringLiteral("reasoning");
|
|
}
|
|
|
|
LlmSegment::Type segmentTypeFromName(const QString& name) {
|
|
if (name == QLatin1String("tool_call"))
|
|
return LlmSegment::Type::ToolCall;
|
|
if (name == QLatin1String("content"))
|
|
return LlmSegment::Type::Content;
|
|
return LlmSegment::Type::Reasoning;
|
|
}
|
|
|
|
// A null QString binds as SQL NULL, which violates the NOT NULL columns;
|
|
// DEFAULT only applies to omitted columns, not explicit NULLs.
|
|
QString sqlText(const QString& value) {
|
|
if (value.isNull())
|
|
return QStringLiteral("");
|
|
return value;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
ChatStore::ChatStore(QObject* parent)
|
|
: QObject(parent), m_connectionName(QUuid::createUuid().toString()) {
|
|
openDb();
|
|
load();
|
|
}
|
|
|
|
ChatStore::~ChatStore() {
|
|
if (m_connectionName.isEmpty())
|
|
return;
|
|
QSqlDatabase db = QSqlDatabase::database(m_connectionName, false);
|
|
db.close();
|
|
QSqlDatabase::removeDatabase(m_connectionName);
|
|
}
|
|
|
|
QSqlDatabase ChatStore::db() const {
|
|
return QSqlDatabase::database(m_connectionName);
|
|
}
|
|
|
|
void ChatStore::openDb() {
|
|
const QString path =
|
|
QStandardPaths::writableLocation(QStandardPaths::GenericStateLocation) +
|
|
QStringLiteral("/zshell/chats.sqlite");
|
|
QDir().mkpath(QFileInfo(path).absolutePath());
|
|
|
|
QSqlDatabase db =
|
|
QSqlDatabase::addDatabase(QStringLiteral("QSQLITE"), m_connectionName);
|
|
db.setDatabaseName(path);
|
|
if (!db.open()) {
|
|
qWarning() << "ChatStore: failed to open database" << path << ":"
|
|
<< db.lastError().text();
|
|
return;
|
|
}
|
|
{
|
|
QSqlQuery pragma(db);
|
|
pragma.exec(QStringLiteral("PRAGMA foreign_keys = ON"));
|
|
}
|
|
{
|
|
QSqlQuery query(db);
|
|
query.exec(
|
|
QStringLiteral(
|
|
"CREATE TABLE IF NOT EXISTS sessions (\n"
|
|
" id TEXT PRIMARY KEY,\n"
|
|
" title TEXT NOT NULL DEFAULT '',\n"
|
|
" icon TEXT NOT NULL DEFAULT '',\n"
|
|
" created_at INTEGER NOT NULL,\n"
|
|
" updated_at INTEGER NOT NULL,\n"
|
|
" pinned INTEGER NOT NULL DEFAULT 0\n"
|
|
")"));
|
|
}
|
|
{
|
|
QSqlQuery query(db);
|
|
query.exec(
|
|
QStringLiteral(
|
|
"CREATE TABLE IF NOT EXISTS messages (\n"
|
|
" id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
|
|
" session_id TEXT NOT NULL REFERENCES sessions (id) "
|
|
"ON DELETE CASCADE,\n"
|
|
" role TEXT NOT NULL,\n"
|
|
" timestamp INTEGER NOT NULL\n"
|
|
")"));
|
|
}
|
|
{
|
|
QSqlQuery query(db);
|
|
query.exec(
|
|
QStringLiteral(
|
|
"CREATE TABLE IF NOT EXISTS generations (\n"
|
|
" id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
|
|
" message_id INTEGER NOT NULL REFERENCES messages "
|
|
"(id) ON DELETE CASCADE,\n"
|
|
" timestamp INTEGER NOT NULL,\n"
|
|
" is_active INTEGER NOT NULL DEFAULT 1\n"
|
|
")"));
|
|
}
|
|
{
|
|
QSqlQuery query(db);
|
|
query.exec(
|
|
QStringLiteral(
|
|
"CREATE TABLE IF NOT EXISTS segments (\n"
|
|
" id INTEGER PRIMARY KEY AUTOINCREMENT,\n"
|
|
" generation_id INTEGER NOT NULL REFERENCES generations "
|
|
"(id) ON DELETE CASCADE,\n"
|
|
" type TEXT NOT NULL,\n"
|
|
" text TEXT NOT NULL DEFAULT '',\n"
|
|
" name TEXT NOT NULL DEFAULT '',\n"
|
|
" tool_call_id TEXT NOT NULL DEFAULT '',\n"
|
|
" arguments TEXT NOT NULL DEFAULT '',\n"
|
|
" result TEXT NOT NULL DEFAULT '',\n"
|
|
" status INTEGER NOT NULL DEFAULT 0,\n"
|
|
" elapsed_ms INTEGER NOT NULL DEFAULT 0,\n"
|
|
" timestamp INTEGER NOT NULL\n"
|
|
")"));
|
|
}
|
|
{
|
|
QSqlQuery query(db);
|
|
query.exec(QStringLiteral(
|
|
"CREATE INDEX IF NOT EXISTS idx_messages_session "
|
|
"ON messages (session_id)"));
|
|
query.exec(QStringLiteral(
|
|
"CREATE INDEX IF NOT EXISTS idx_generations_message "
|
|
"ON generations (message_id)"));
|
|
query.exec(QStringLiteral(
|
|
"CREATE INDEX IF NOT EXISTS idx_segments_generation "
|
|
"ON segments (generation_id)"));
|
|
}
|
|
}
|
|
|
|
int ChatStore::count() const {
|
|
return static_cast<int>(m_sessions.size());
|
|
}
|
|
|
|
QVariantList ChatStore::values() const {
|
|
QVariantList vals;
|
|
vals.reserve(m_sessions.size());
|
|
for (const auto* session : m_sessions)
|
|
vals.append(QVariant::fromValue(session));
|
|
return vals;
|
|
}
|
|
|
|
ChatSession* ChatStore::at(int index) const {
|
|
if (index < 0 || index >= m_sessions.size())
|
|
return nullptr;
|
|
return m_sessions.at(index);
|
|
}
|
|
|
|
ChatSession* ChatStore::insert(int index) {
|
|
const qint64 now = QDateTime::currentMSecsSinceEpoch();
|
|
const QString id = QUuid::createUuid().toString();
|
|
{
|
|
QSqlQuery query(db());
|
|
query.prepare(
|
|
"INSERT INTO sessions (id, title, created_at, updated_at) "
|
|
"VALUES (:id, '', :created_at, :updated_at)");
|
|
query.bindValue(":id", id);
|
|
query.bindValue(":created_at", now);
|
|
query.bindValue(":updated_at", now);
|
|
if (!query.exec())
|
|
qWarning() << "ChatStore: failed to insert session" << id << ":"
|
|
<< query.lastError().text();
|
|
}
|
|
auto* session = new ChatSession(id, this);
|
|
session->setMeta(QString(), now, now, 0);
|
|
const int pos = index >= 0 && index <= m_sessions.size() ? index : 0;
|
|
m_sessions.insert(pos, session);
|
|
Q_EMIT countChanged();
|
|
Q_EMIT valuesChanged();
|
|
return session;
|
|
}
|
|
|
|
void ChatStore::remove(int index) {
|
|
removeSession(at(index));
|
|
}
|
|
|
|
void ChatStore::remove(ChatSession* chat) {
|
|
removeSession(chat);
|
|
}
|
|
|
|
void ChatStore::removeSession(ChatSession* session) {
|
|
if (!session || !m_sessions.contains(session))
|
|
return;
|
|
const QList<ChatSession*> before = m_sessions;
|
|
Q_EMIT sessionRemoved(session);
|
|
{
|
|
QSqlQuery query(db());
|
|
query.prepare("DELETE FROM sessions WHERE id = :id");
|
|
query.bindValue(":id", session->id());
|
|
query.exec();
|
|
}
|
|
m_sessions.removeOne(session);
|
|
session->deleteLater();
|
|
notify(before);
|
|
}
|
|
|
|
void ChatStore::move(int from, int to) {
|
|
if (from < 0 || from >= m_sessions.size() || to < 0 ||
|
|
to >= m_sessions.size() || from == to)
|
|
return;
|
|
m_sessions.move(from, to);
|
|
Q_EMIT valuesChanged();
|
|
}
|
|
|
|
void ChatStore::clear() {
|
|
const QList<ChatSession*> sessions = m_sessions;
|
|
for (ChatSession* session : sessions)
|
|
removeSession(session);
|
|
}
|
|
|
|
ChatSession* ChatStore::sessionById(const QString& id) {
|
|
for (auto* session : m_sessions)
|
|
if (session->id() == id)
|
|
return session;
|
|
return nullptr;
|
|
}
|
|
|
|
void ChatStore::setLlmClient(LlmClient* client) {
|
|
m_llmClient = client;
|
|
}
|
|
|
|
void ChatStore::persist(ChatSession* session) {
|
|
if (!session || !m_sessions.contains(session))
|
|
return;
|
|
session->setUpdatedAt(QDateTime::currentMSecsSinceEpoch());
|
|
session->ensureLoaded();
|
|
if (!saveSession(session))
|
|
return;
|
|
sortAndNotify();
|
|
}
|
|
|
|
void ChatStore::saveMeta(ChatSession* session) {
|
|
if (!session || !m_sessions.contains(session))
|
|
return;
|
|
QSqlQuery query(db());
|
|
query.prepare("UPDATE sessions SET title = :title, icon = :icon "
|
|
"WHERE id = :id");
|
|
query.bindValue(":title", sqlText(session->title()));
|
|
query.bindValue(":icon", sqlText(session->icon()));
|
|
query.bindValue(":id", session->id());
|
|
if (!query.exec())
|
|
qWarning() << "ChatStore: failed to save meta for" << session->id()
|
|
<< ":" << query.lastError().text();
|
|
}
|
|
|
|
bool ChatStore::saveSession(ChatSession* session) {
|
|
const QString id = session->id();
|
|
QSqlDatabase handle = db();
|
|
if (!handle.transaction()) {
|
|
qWarning() << "ChatStore: failed to begin transaction:"
|
|
<< handle.lastError().text();
|
|
return false;
|
|
}
|
|
bool ok = true;
|
|
{
|
|
QSqlQuery query(handle);
|
|
query.prepare(
|
|
"UPDATE sessions SET title = :title, updated_at = :updated_at "
|
|
"WHERE id = :id");
|
|
query.bindValue(":title", sqlText(session->title()));
|
|
query.bindValue(":updated_at", session->updatedAtMs());
|
|
query.bindValue(":id", session->id());
|
|
ok = query.exec();
|
|
}
|
|
if (ok) {
|
|
QSqlQuery query(handle);
|
|
query.prepare("DELETE FROM messages WHERE session_id = :id");
|
|
query.bindValue(":id", session->id());
|
|
ok = query.exec();
|
|
}
|
|
if (ok) {
|
|
QSqlQuery messageInsert(handle);
|
|
ok = messageInsert.prepare(
|
|
"INSERT INTO messages (session_id, role, timestamp) "
|
|
"VALUES (:id, :role, :timestamp)");
|
|
QSqlQuery generationInsert(handle);
|
|
ok = ok && generationInsert.prepare(
|
|
"INSERT INTO generations (message_id, timestamp, is_active) "
|
|
"VALUES (:mid, :timestamp, :is_active)");
|
|
QSqlQuery segmentInsert(handle);
|
|
ok = ok && segmentInsert.prepare(
|
|
"INSERT INTO segments (generation_id, type, text, name, "
|
|
"tool_call_id, arguments, result, status, elapsed_ms, "
|
|
"timestamp) VALUES (:gid, :type, :text, :name, "
|
|
":tool_call_id, :arguments, :result, :status, :elapsed_ms, "
|
|
":timestamp)");
|
|
// The model holds messages most recent first; the database keeps
|
|
// natural rowid order, so iterate from the oldest row up.
|
|
const auto* model = session->messagesModel();
|
|
for (int row = model->rowCount() - 1; ok && row >= 0; --row) {
|
|
const auto* message = model->at(row);
|
|
messageInsert.bindValue(":id", session->id());
|
|
messageInsert.bindValue(
|
|
":role",
|
|
message->role() == ChatMessage::Role::User
|
|
? QStringLiteral("user")
|
|
: QStringLiteral("assistant"));
|
|
messageInsert.bindValue(":timestamp", message->timestamp());
|
|
if (!messageInsert.exec()) {
|
|
ok = false;
|
|
qWarning() << "ChatStore: saveSession" << id
|
|
<< "message insert failed:"
|
|
<< messageInsert.lastError().text();
|
|
break;
|
|
}
|
|
const int messageId = messageInsert.lastInsertId().toInt();
|
|
for (int i = 0; ok && i < message->generationCount(); ++i) {
|
|
const auto* generation = message->generation(i);
|
|
generationInsert.bindValue(":mid", messageId);
|
|
generationInsert.bindValue(
|
|
":timestamp", generation->timestamp());
|
|
generationInsert.bindValue(
|
|
":is_active",
|
|
i == message->activeGenerationIndex() ? 1 : 0);
|
|
if (!generationInsert.exec()) {
|
|
ok = false;
|
|
qWarning() << "ChatStore: saveSession" << id
|
|
<< "generation insert failed:"
|
|
<< generationInsert.lastError().text();
|
|
break;
|
|
}
|
|
const int generationId =
|
|
generationInsert.lastInsertId().toInt();
|
|
for (const auto* segment : generation->segments()) {
|
|
segmentInsert.bindValue(":gid", generationId);
|
|
segmentInsert.bindValue(
|
|
":type", segmentTypeName(segment->type()));
|
|
segmentInsert.bindValue(":text", sqlText(segment->text()));
|
|
segmentInsert.bindValue(":name", sqlText(segment->name()));
|
|
segmentInsert.bindValue(
|
|
":tool_call_id", sqlText(segment->toolCallId()));
|
|
segmentInsert.bindValue(
|
|
":arguments", sqlText(segment->arguments()));
|
|
segmentInsert.bindValue(":result", sqlText(segment->result()));
|
|
segmentInsert.bindValue(
|
|
":status", static_cast<int>(segment->status()));
|
|
segmentInsert.bindValue(
|
|
":elapsed_ms", segment->elapsedMs());
|
|
segmentInsert.bindValue(
|
|
":timestamp", segment->timestamp());
|
|
if (!segmentInsert.exec()) {
|
|
ok = false;
|
|
qWarning() << "ChatStore: saveSession" << id
|
|
<< "segment insert failed:"
|
|
<< segmentInsert.lastError().text();
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (!ok || !handle.commit()) {
|
|
qWarning() << "ChatStore: saveSession" << id << "commit failed, rolling back";
|
|
handle.rollback();
|
|
ok = false;
|
|
qWarning() << "ChatStore: failed to save session" << session->id() << ":"
|
|
<< handle.lastError().text();
|
|
}
|
|
return ok;
|
|
}
|
|
|
|
void ChatStore::loadMessagesInto(ChatSession* session) {
|
|
// Newest first so the model receives rows in display order.
|
|
QSqlQuery query(db());
|
|
query.prepare(
|
|
"SELECT id, role, timestamp FROM messages WHERE session_id = :id "
|
|
"ORDER BY rowid DESC");
|
|
query.bindValue(":id", session->id());
|
|
if (!query.exec()) {
|
|
qWarning() << "ChatStore: failed to load messages for" << session->id()
|
|
<< ":" << query.lastError().text();
|
|
return;
|
|
}
|
|
auto* model = session->messagesModel();
|
|
QList<ChatMessage*> messages;
|
|
while (query.next()) {
|
|
const int messageId = query.value(0).toInt();
|
|
auto* message = model->createMessage(
|
|
query.value(1).toString() == QLatin1String("user")
|
|
? ChatMessage::Role::User
|
|
: ChatMessage::Role::Assistant,
|
|
query.value(2).toLongLong());
|
|
QSqlQuery generationQuery(db());
|
|
generationQuery.prepare(
|
|
"SELECT id, timestamp, is_active FROM generations "
|
|
"WHERE message_id = :mid ORDER BY rowid");
|
|
generationQuery.bindValue(":mid", messageId);
|
|
int activeIndex = 0;
|
|
if (generationQuery.exec()) {
|
|
int index = 0;
|
|
while (generationQuery.next()) {
|
|
auto* generation = message->addGeneration(
|
|
generationQuery.value(1).toLongLong());
|
|
QSqlQuery segmentQuery(db());
|
|
segmentQuery.prepare(
|
|
"SELECT type, text, name, tool_call_id, arguments, "
|
|
"result, status, elapsed_ms, timestamp FROM segments "
|
|
"WHERE generation_id = :gid ORDER BY rowid");
|
|
segmentQuery.bindValue(
|
|
":gid", generationQuery.value(0).toInt());
|
|
if (segmentQuery.exec()) {
|
|
while (segmentQuery.next()) {
|
|
auto* segment = new LlmSegment(
|
|
segmentTypeFromName(
|
|
segmentQuery.value(0).toString()),
|
|
segmentQuery.value(8).toLongLong(),
|
|
generation);
|
|
segment->setText(
|
|
segmentQuery.value(1).toString());
|
|
segment->setName(
|
|
segmentQuery.value(2).toString());
|
|
segment->setToolCallId(
|
|
segmentQuery.value(3).toString());
|
|
segment->appendArguments(
|
|
segmentQuery.value(4).toString());
|
|
segment->setResult(
|
|
segmentQuery.value(5).toString());
|
|
segment->setStatus(
|
|
static_cast<LlmSegment::Status>(
|
|
segmentQuery.value(6).toInt()));
|
|
segment->restore(
|
|
segmentQuery.value(7).toLongLong());
|
|
generation->addSegment(segment);
|
|
}
|
|
} else {
|
|
qWarning() << "ChatStore: failed to load segments for "
|
|
<< "generation"
|
|
<< generationQuery.value(0).toInt()
|
|
<< ":"
|
|
<< segmentQuery.lastError().text();
|
|
}
|
|
if (generationQuery.value(2).toInt() != 0)
|
|
activeIndex = index;
|
|
++index;
|
|
}
|
|
} else {
|
|
qWarning() << "ChatStore: failed to load generations for message"
|
|
<< messageId << ":"
|
|
<< generationQuery.lastError().text();
|
|
}
|
|
message->setActiveGeneration(activeIndex);
|
|
messages.append(message);
|
|
}
|
|
session->adoptMessages(messages);
|
|
}
|
|
|
|
void ChatStore::load() {
|
|
QSqlQuery query(db());
|
|
query.exec(
|
|
"SELECT s.id, s.title, s.icon, s.created_at, s.updated_at, s.pinned, "
|
|
"(SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) AS cnt "
|
|
"FROM sessions s ORDER BY s.pinned DESC, s.updated_at DESC");
|
|
while (query.next()) {
|
|
auto* session = new ChatSession(query.value(0).toString(), this);
|
|
session->setMeta(
|
|
query.value(1).toString(),
|
|
query.value(3).toLongLong(),
|
|
query.value(4).toLongLong(),
|
|
query.value(6).toInt());
|
|
session->setIcon(query.value(2).toString());
|
|
session->setPinned(query.value(5).toInt() != 0);
|
|
m_sessions.append(session);
|
|
}
|
|
sortAndNotify();
|
|
}
|
|
|
|
void ChatStore::sortAndNotify() {
|
|
const QList<ChatSession*> before = m_sessions;
|
|
std::stable_sort(
|
|
m_sessions.begin(),
|
|
m_sessions.end(),
|
|
[](const ChatSession* a, const ChatSession* b) {
|
|
if (a->pinned() != b->pinned())
|
|
return a->pinned() > b->pinned();
|
|
return a->updatedAtMs() > b->updatedAtMs();
|
|
});
|
|
notify(before);
|
|
}
|
|
|
|
void ChatStore::notify(const QList<ChatSession*>& before) {
|
|
if (before.size() != m_sessions.size())
|
|
Q_EMIT countChanged();
|
|
bool same = before.size() == m_sessions.size();
|
|
for (int i = 0; same && i < m_sessions.size(); ++i)
|
|
if (before.at(i) != m_sessions.at(i))
|
|
same = false;
|
|
if (!same)
|
|
Q_EMIT valuesChanged();
|
|
}
|
|
|
|
} // namespace ZShell::llm
|