mirror of
https://github.com/nomic-ai/gpt4all
synced 2024-11-08 07:10:32 +00:00
2177 lines
69 KiB
C++
2177 lines
69 KiB
C++
#include "database.h"
|
|
|
|
#include "mysettings.h"
|
|
|
|
#include <usearch/index_plugins.hpp>
|
|
|
|
#include <QDebug>
|
|
#include <QDir>
|
|
#include <QDirIterator>
|
|
#include <QElapsedTimer>
|
|
#include <QFile>
|
|
#include <QFileSystemWatcher>
|
|
#include <QIODevice>
|
|
#include <QPdfDocument>
|
|
#include <QPdfSelection>
|
|
#include <QRegularExpression>
|
|
#include <QSqlError>
|
|
#include <QSqlQuery>
|
|
#include <QTextStream>
|
|
#include <QTimer>
|
|
#include <QVariant>
|
|
#include <Qt>
|
|
#include <QtGlobal>
|
|
#include <QtLogging>
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <optional>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
using namespace Qt::Literals::StringLiterals;
|
|
namespace us = unum::usearch;
|
|
|
|
//#define DEBUG
|
|
//#define DEBUG_EXAMPLE
|
|
|
|
namespace {
|
|
|
|
/* QFile that checks input for binary data. If seen, it fails the read and returns true
|
|
* for binarySeen(). */
|
|
class BinaryDetectingFile: public QFile {
|
|
public:
|
|
using QFile::QFile;
|
|
|
|
bool binarySeen() const { return m_binarySeen; }
|
|
|
|
protected:
|
|
qint64 readData(char *data, qint64 maxSize) override {
|
|
qint64 res = QFile::readData(data, maxSize);
|
|
return checkData(data, res);
|
|
}
|
|
|
|
qint64 readLineData(char *data, qint64 maxSize) override {
|
|
qint64 res = QFile::readLineData(data, maxSize);
|
|
return checkData(data, res);
|
|
}
|
|
|
|
private:
|
|
qint64 checkData(const char *data, qint64 size) {
|
|
Q_ASSERT(!isTextModeEnabled()); // We need raw bytes from the underlying QFile
|
|
if (size != -1 && !m_binarySeen) {
|
|
for (qint64 i = 0; i < size; i++) {
|
|
/* Control characters we should never see in plain text:
|
|
* 0x00 NUL - 0x06 ACK
|
|
* 0x0E SO - 0x1A SUB
|
|
* 0x1C FS - 0x1F US */
|
|
auto c = static_cast<unsigned char>(data[i]);
|
|
if (c < 0x07 || (c >= 0x0E && c < 0x1B) || (c >= 0x1C && c < 0x20)) {
|
|
m_binarySeen = true;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
return m_binarySeen ? -1 : size;
|
|
}
|
|
|
|
bool m_binarySeen = false;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
static int s_batchSize = 100;
|
|
|
|
static const QString INIT_DB_SQL[] = {
|
|
// automatically free unused disk space
|
|
u"pragma auto_vacuum = FULL;"_s,
|
|
// create tables
|
|
uR"(
|
|
create table chunks(
|
|
id integer primary key autoincrement,
|
|
document_id integer not null,
|
|
chunk_text text not null,
|
|
file text not null,
|
|
title text,
|
|
author text,
|
|
subject text,
|
|
keywords text,
|
|
page integer,
|
|
line_from integer,
|
|
line_to integer,
|
|
words integer default 0 not null,
|
|
tokens integer default 0 not null,
|
|
foreign key(document_id) references documents(id)
|
|
);
|
|
)"_s, uR"(
|
|
create table collections(
|
|
id integer primary key,
|
|
name text unique not null,
|
|
start_update_time integer,
|
|
last_update_time integer,
|
|
embedding_model text
|
|
);
|
|
)"_s, uR"(
|
|
create table folders(
|
|
id integer primary key autoincrement,
|
|
path text unique not null
|
|
);
|
|
)"_s, uR"(
|
|
create table collection_items(
|
|
collection_id integer not null,
|
|
folder_id integer not null,
|
|
foreign key(collection_id) references collections(id)
|
|
foreign key(folder_id) references folders(id),
|
|
unique(collection_id, folder_id)
|
|
);
|
|
)"_s, uR"(
|
|
create table documents(
|
|
id integer primary key,
|
|
folder_id integer not null,
|
|
document_time integer not null,
|
|
document_path text unique not null,
|
|
foreign key(folder_id) references folders(id)
|
|
);
|
|
)"_s, uR"(
|
|
create table embeddings(
|
|
model text not null,
|
|
folder_id integer not null,
|
|
chunk_id integer not null,
|
|
embedding blob not null,
|
|
primary key(model, folder_id, chunk_id),
|
|
foreign key(folder_id) references folders(id),
|
|
foreign key(chunk_id) references chunks(id),
|
|
unique(model, chunk_id)
|
|
);
|
|
)"_s,
|
|
};
|
|
|
|
static const QString INSERT_CHUNK_SQL = uR"(
|
|
insert into chunks(document_id, chunk_text,
|
|
file, title, author, subject, keywords, page, line_from, line_to, words)
|
|
values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
returning id;
|
|
)"_s;
|
|
|
|
static const QString DELETE_CHUNKS_SQL[] = {
|
|
uR"(
|
|
delete from embeddings
|
|
where chunk_id in (
|
|
select id from chunks where document_id = ?
|
|
);
|
|
)"_s, uR"(
|
|
delete from chunks where document_id = ?;
|
|
)"_s,
|
|
};
|
|
|
|
static const QString SELECT_CHUNKS_BY_DOCUMENT_SQL = uR"(
|
|
select id from chunks WHERE document_id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_CHUNKS_SQL = uR"(
|
|
select c.id, d.document_time, d.document_path, c.chunk_text, c.file, c.title, c.author, c.page, c.line_from, c.line_to, co.name
|
|
from chunks c
|
|
join documents d on d.id = c.document_id
|
|
join folders f on f.id = d.folder_id
|
|
join collection_items ci on ci.folder_id = f.id
|
|
join collections co on co.id = ci.collection_id
|
|
where c.id in (%1);
|
|
)"_s;
|
|
|
|
static const QString SELECT_UNCOMPLETED_CHUNKS_SQL = uR"(
|
|
select co.name, co.embedding_model, c.id, d.folder_id, c.chunk_text
|
|
from chunks c
|
|
join documents d on d.id = c.document_id
|
|
join folders f on f.id = d.folder_id
|
|
join collection_items ci on ci.folder_id = f.id
|
|
join collections co on co.id = ci.collection_id and co.embedding_model is not null
|
|
where not exists(
|
|
select 1
|
|
from embeddings e
|
|
where e.chunk_id = c.id and e.model = co.embedding_model
|
|
);
|
|
)"_s;
|
|
|
|
static const QString SELECT_COUNT_CHUNKS_SQL = uR"(
|
|
select count(c.id)
|
|
from chunks c
|
|
join documents d on d.id = c.document_id
|
|
where d.folder_id = ?;
|
|
)"_s;
|
|
|
|
static bool addChunk(QSqlQuery &q, int document_id, const QString &chunk_text, const QString &file,
|
|
const QString &title, const QString &author, const QString &subject, const QString &keywords,
|
|
int page, int from, int to, int words, int *chunk_id)
|
|
{
|
|
if (!q.prepare(INSERT_CHUNK_SQL))
|
|
return false;
|
|
q.addBindValue(document_id);
|
|
q.addBindValue(chunk_text);
|
|
q.addBindValue(file);
|
|
q.addBindValue(title);
|
|
q.addBindValue(author);
|
|
q.addBindValue(subject);
|
|
q.addBindValue(keywords);
|
|
q.addBindValue(page);
|
|
q.addBindValue(from);
|
|
q.addBindValue(to);
|
|
q.addBindValue(words);
|
|
if (!q.exec() || !q.next())
|
|
return false;
|
|
*chunk_id = q.value(0).toInt();
|
|
return true;
|
|
}
|
|
|
|
static bool removeChunksByDocumentId(QSqlQuery &q, int document_id)
|
|
{
|
|
for (const auto &cmd: DELETE_CHUNKS_SQL) {
|
|
if (!q.prepare(cmd))
|
|
return false;
|
|
q.addBindValue(document_id);
|
|
if (!q.exec())
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
#define NAMED_PAIR(name, typea, a, typeb, b) \
|
|
struct name { typea a; typeb b; }; \
|
|
static bool operator==(const name &x, const name &y) { return x.a == y.a && x.b == y.b; } \
|
|
static size_t qHash(const name &x, size_t seed) { return qHashMulti(seed, x.a, x.b); }
|
|
|
|
// struct compared by embedding key, can be extended with additional unique data
|
|
NAMED_PAIR(EmbeddingKey, QString, embedding_model, int, chunk_id)
|
|
|
|
namespace {
|
|
struct IncompleteChunk: EmbeddingKey { int folder_id; QString text; };
|
|
} // namespace
|
|
|
|
static bool selectAllUncompletedChunks(QSqlQuery &q, QHash<IncompleteChunk, QStringList> &chunks)
|
|
{
|
|
if (!q.exec(SELECT_UNCOMPLETED_CHUNKS_SQL))
|
|
return false;
|
|
while (q.next()) {
|
|
QString collection = q.value(0).toString();
|
|
IncompleteChunk ic {
|
|
/*embedding_model =*/ q.value(1).toString(),
|
|
/*chunk_id =*/ q.value(2).toInt(),
|
|
/*folder_id =*/ q.value(3).toInt(),
|
|
/*text =*/ q.value(4).toString(),
|
|
};
|
|
chunks[ic] << collection;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static bool selectCountChunks(QSqlQuery &q, int folder_id, int &count)
|
|
{
|
|
if (!q.prepare(SELECT_COUNT_CHUNKS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_id);
|
|
if (!q.exec())
|
|
return false;
|
|
if (!q.next()) {
|
|
count = 0;
|
|
return false;
|
|
}
|
|
count = q.value(0).toInt();
|
|
return true;
|
|
}
|
|
|
|
static bool selectChunk(QSqlQuery &q, const QList<int> &chunk_ids, int retrievalSize)
|
|
{
|
|
QString chunk_ids_str = QString::number(chunk_ids[0]);
|
|
for (size_t i = 1; i < chunk_ids.size(); ++i)
|
|
chunk_ids_str += "," + QString::number(chunk_ids[i]);
|
|
const QString formatted_query = SELECT_CHUNKS_SQL.arg(chunk_ids_str);
|
|
if (!q.prepare(formatted_query))
|
|
return false;
|
|
return q.exec();
|
|
}
|
|
|
|
static const QString INSERT_COLLECTION_SQL = uR"(
|
|
insert into collections(name, start_update_time, last_update_time, embedding_model)
|
|
values(?, ?, ?, ?)
|
|
returning id;
|
|
)"_s;
|
|
|
|
static const QString DELETE_COLLECTION_SQL = uR"(
|
|
delete from collections where name = ? and folder_id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_FOLDERS_FROM_COLLECTIONS_SQL = uR"(
|
|
select f.id, f.path
|
|
from collections c
|
|
join collection_items ci on ci.collection_id = c.id
|
|
join folders f on ci.folder_id = f.id
|
|
where c.name = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_COLLECTIONS_SQL_V1 = uR"(
|
|
select c.collection_name, f.folder_path, f.id
|
|
from collections c
|
|
join folders f on c.folder_id = f.id
|
|
order by c.collection_name asc, f.folder_path asc;
|
|
)"_s;
|
|
|
|
static const QString SELECT_COLLECTIONS_SQL_V2 = uR"(
|
|
select c.id, c.name, f.path, f.id, c.start_update_time, c.last_update_time, c.embedding_model
|
|
from collections c
|
|
join collection_items ci on ci.collection_id = c.id
|
|
join folders f on ci.folder_id = f.id
|
|
order by c.name asc, f.path asc;
|
|
)"_s;
|
|
|
|
static const QString SELECT_COLLECTION_BY_NAME_SQL = uR"(
|
|
select id, name, start_update_time, last_update_time, embedding_model
|
|
from collections c
|
|
where name = ?;
|
|
)"_s;
|
|
|
|
static const QString SET_COLLECTION_EMBEDDING_MODEL_SQL = uR"(
|
|
update collections
|
|
set embedding_model = ?
|
|
where name = ?;
|
|
)"_s;
|
|
|
|
static const QString UPDATE_START_UPDATE_TIME_SQL = uR"(
|
|
update collections set start_update_time = ? where id = ?;
|
|
)"_s;
|
|
|
|
static const QString UPDATE_LAST_UPDATE_TIME_SQL = uR"(
|
|
update collections set last_update_time = ? where id = ?;
|
|
)"_s;
|
|
|
|
static bool addCollection(QSqlQuery &q, const QString &collection_name, const QDateTime &start_update,
|
|
const QDateTime &last_update, const QString &embedding_model, CollectionItem &item)
|
|
{
|
|
if (!q.prepare(INSERT_COLLECTION_SQL))
|
|
return false;
|
|
q.addBindValue(collection_name);
|
|
q.addBindValue(start_update);
|
|
q.addBindValue(last_update);
|
|
q.addBindValue(embedding_model);
|
|
if (!q.exec() || !q.next())
|
|
return false;
|
|
item.collection_id = q.value(0).toInt();
|
|
item.collection = collection_name;
|
|
item.embeddingModel = embedding_model;
|
|
return true;
|
|
}
|
|
|
|
static bool removeCollection(QSqlQuery &q, const QString &collection_name, int folder_id)
|
|
{
|
|
if (!q.prepare(DELETE_COLLECTION_SQL))
|
|
return false;
|
|
q.addBindValue(collection_name);
|
|
q.addBindValue(folder_id);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool selectFoldersFromCollection(QSqlQuery &q, const QString &collection_name, QList<QPair<int, QString>> *folders)
|
|
{
|
|
if (!q.prepare(SELECT_FOLDERS_FROM_COLLECTIONS_SQL))
|
|
return false;
|
|
q.addBindValue(collection_name);
|
|
if (!q.exec())
|
|
return false;
|
|
while (q.next())
|
|
folders->append({q.value(0).toInt(), q.value(1).toString()});
|
|
return true;
|
|
}
|
|
|
|
static QList<CollectionItem> sqlExtractCollections(QSqlQuery &q, bool with_folder = false, int version = LOCALDOCS_VERSION)
|
|
{
|
|
QList<CollectionItem> collections;
|
|
while (q.next()) {
|
|
CollectionItem i;
|
|
int idx = 0;
|
|
if (version >= 2)
|
|
i.collection_id = q.value(idx++).toInt();
|
|
i.collection = q.value(idx++).toString();
|
|
if (with_folder) {
|
|
i.folder_path = q.value(idx++).toString();
|
|
i.folder_id = q.value(idx++).toInt();
|
|
}
|
|
i.indexing = false;
|
|
i.installed = true;
|
|
|
|
if (version >= 2) {
|
|
bool ok;
|
|
const qint64 start_update = q.value(idx++).toLongLong(&ok);
|
|
if (ok) i.startUpdate = QDateTime::fromMSecsSinceEpoch(start_update);
|
|
const qint64 last_update = q.value(idx++).toLongLong(&ok);
|
|
if (ok) i.lastUpdate = QDateTime::fromMSecsSinceEpoch(last_update);
|
|
|
|
i.embeddingModel = q.value(idx++).toString();
|
|
}
|
|
if (i.embeddingModel.isNull()) {
|
|
// unknown embedding model -> need to re-index
|
|
i.forceIndexing = true;
|
|
}
|
|
|
|
collections << i;
|
|
}
|
|
return collections;
|
|
}
|
|
|
|
static bool selectAllFromCollections(QSqlQuery &q, QList<CollectionItem> *collections, int version = LOCALDOCS_VERSION)
|
|
{
|
|
|
|
switch (version) {
|
|
case 1:
|
|
if (!q.prepare(SELECT_COLLECTIONS_SQL_V1))
|
|
return false;
|
|
break;
|
|
case 2:
|
|
if (!q.prepare(SELECT_COLLECTIONS_SQL_V2))
|
|
return false;
|
|
break;
|
|
default:
|
|
Q_UNREACHABLE();
|
|
return false;
|
|
}
|
|
|
|
if (!q.exec())
|
|
return false;
|
|
*collections = sqlExtractCollections(q, true, version);
|
|
return true;
|
|
}
|
|
|
|
static bool selectCollectionByName(QSqlQuery &q, const QString &name, std::optional<CollectionItem> &collection)
|
|
{
|
|
if (!q.prepare(SELECT_COLLECTION_BY_NAME_SQL))
|
|
return false;
|
|
q.addBindValue(name);
|
|
if (!q.exec())
|
|
return false;
|
|
QList<CollectionItem> collections = sqlExtractCollections(q);
|
|
Q_ASSERT(collections.count() <= 1);
|
|
collection.reset();
|
|
if (!collections.isEmpty())
|
|
collection = collections.first();
|
|
return true;
|
|
}
|
|
|
|
static bool setCollectionEmbeddingModel(QSqlQuery &q, const QString &collection_name, const QString &embedding_model)
|
|
{
|
|
if (!q.prepare(SET_COLLECTION_EMBEDDING_MODEL_SQL))
|
|
return false;
|
|
q.addBindValue(embedding_model);
|
|
q.addBindValue(collection_name);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool updateStartUpdateTime(QSqlQuery &q, int id, qint64 update_time)
|
|
{
|
|
if (!q.prepare(UPDATE_START_UPDATE_TIME_SQL))
|
|
return false;
|
|
q.addBindValue(update_time);
|
|
q.addBindValue(id);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool updateLastUpdateTime(QSqlQuery &q, int id, qint64 update_time)
|
|
{
|
|
if (!q.prepare(UPDATE_LAST_UPDATE_TIME_SQL))
|
|
return false;
|
|
q.addBindValue(update_time);
|
|
q.addBindValue(id);
|
|
return q.exec();
|
|
}
|
|
|
|
static const QString INSERT_FOLDERS_SQL = uR"(
|
|
insert into folders(path) values(?);
|
|
)"_s;
|
|
|
|
static const QString DELETE_FOLDERS_SQL = uR"(
|
|
delete from folders where id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_FOLDERS_FROM_PATH_SQL = uR"(
|
|
select id from folders where path = ?;
|
|
)"_s;
|
|
|
|
static const QString GET_FOLDER_EMBEDDING_MODEL_SQL = uR"(
|
|
select co.embedding_model
|
|
from collections co
|
|
join collection_items ci on ci.collection_id = co.id
|
|
where ci.folder_id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_ALL_FOLDERPATHS_SQL = uR"(
|
|
select path from folders;
|
|
)"_s;
|
|
|
|
static const QString FOLDER_REMOVE_ALL_DOCS_SQL[] = {
|
|
uR"(
|
|
delete from embeddings
|
|
where chunk_id in (
|
|
select c.id
|
|
from chunks c
|
|
join documents d on d.id = c.document_id
|
|
join folders f on f.id = d.folder_id
|
|
where f.path = ?
|
|
);
|
|
)"_s, uR"(
|
|
delete from chunks
|
|
where document_id in (
|
|
select d.id
|
|
from documents d
|
|
join folders f on f.id = d.folder_id
|
|
where f.path = ?
|
|
);
|
|
)"_s, uR"(
|
|
delete from documents
|
|
where id in (
|
|
select d.id
|
|
from documents d
|
|
join folders f on f.id = d.folder_id
|
|
where f.path = ?
|
|
);
|
|
)"_s,
|
|
};
|
|
|
|
static bool addFolderToDB(QSqlQuery &q, const QString &folder_path, int *folder_id)
|
|
{
|
|
if (!q.prepare(INSERT_FOLDERS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_path);
|
|
if (!q.exec())
|
|
return false;
|
|
*folder_id = q.lastInsertId().toInt();
|
|
return true;
|
|
}
|
|
|
|
static bool removeFolderFromDB(QSqlQuery &q, int folder_id)
|
|
{
|
|
if (!q.prepare(DELETE_FOLDERS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_id);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool selectFolder(QSqlQuery &q, const QString &folder_path, int *id)
|
|
{
|
|
if (!q.prepare(SELECT_FOLDERS_FROM_PATH_SQL))
|
|
return false;
|
|
q.addBindValue(folder_path);
|
|
if (!q.exec())
|
|
return false;
|
|
Q_ASSERT(q.size() < 2);
|
|
if (q.next())
|
|
*id = q.value(0).toInt();
|
|
return true;
|
|
}
|
|
|
|
static bool sqlGetFolderEmbeddingModel(QSqlQuery &q, int id, QString &embedding_model)
|
|
{
|
|
if (!q.prepare(GET_FOLDER_EMBEDDING_MODEL_SQL))
|
|
return false;
|
|
q.addBindValue(id);
|
|
if (!q.exec() || !q.next())
|
|
return false;
|
|
// FIXME(jared): there may be more than one if a folder is shared between collections
|
|
Q_ASSERT(q.size() < 2);
|
|
embedding_model = q.value(0).toString();
|
|
return true;
|
|
}
|
|
|
|
static bool selectAllFolderPaths(QSqlQuery &q, QList<QString> *folder_paths)
|
|
{
|
|
if (!q.prepare(SELECT_ALL_FOLDERPATHS_SQL))
|
|
return false;
|
|
if (!q.exec())
|
|
return false;
|
|
while (q.next())
|
|
folder_paths->append(q.value(0).toString());
|
|
return true;
|
|
}
|
|
|
|
static bool sqlRemoveDocsByFolderPath(QSqlQuery &q, const QString &path)
|
|
{
|
|
for (const auto &cmd: FOLDER_REMOVE_ALL_DOCS_SQL) {
|
|
if (!q.prepare(cmd))
|
|
return false;
|
|
q.addBindValue(path);
|
|
if (!q.exec())
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static const QString INSERT_COLLECTION_ITEM_SQL = uR"(
|
|
insert into collection_items(collection_id, folder_id)
|
|
values(?, ?)
|
|
on conflict do nothing;
|
|
)"_s;
|
|
|
|
static const QString DELETE_COLLECTION_FOLDER_SQL = uR"(
|
|
delete from collection_items
|
|
where collection_id = (select id from collections where name = :name) and folder_id = :folder_id
|
|
returning (select count(*) from collection_items where folder_id = :folder_id);
|
|
)"_s;
|
|
|
|
static const QString PRUNE_COLLECTIONS_SQL = uR"(
|
|
delete from collections
|
|
where id not in (select collection_id from collection_items);
|
|
)"_s;
|
|
|
|
// 0 = already exists, 1 = added, -1 = error
|
|
static int addCollectionItem(QSqlQuery &q, int collection_id, int folder_id)
|
|
{
|
|
if (!q.prepare(INSERT_COLLECTION_ITEM_SQL))
|
|
return -1;
|
|
q.addBindValue(collection_id);
|
|
q.addBindValue(folder_id);
|
|
if (q.exec())
|
|
return q.numRowsAffected();
|
|
return -1;
|
|
}
|
|
|
|
// returns the number of remaining references to the folder, or -1 on error
|
|
static int removeCollectionFolder(QSqlQuery &q, const QString &collection_name, int folder_id)
|
|
{
|
|
if (!q.prepare(DELETE_COLLECTION_FOLDER_SQL))
|
|
return -1;
|
|
q.bindValue(":name", collection_name);
|
|
q.bindValue(":folder_id", folder_id);
|
|
if (!q.exec() || !q.next())
|
|
return -1;
|
|
return q.value(0).toInt();
|
|
}
|
|
|
|
static bool sqlPruneCollections(QSqlQuery &q)
|
|
{
|
|
return q.exec(PRUNE_COLLECTIONS_SQL);
|
|
}
|
|
|
|
static const QString INSERT_DOCUMENTS_SQL = uR"(
|
|
insert into documents(folder_id, document_time, document_path) values(?, ?, ?);
|
|
)"_s;
|
|
|
|
static const QString UPDATE_DOCUMENT_TIME_SQL = uR"(
|
|
update documents set document_time = ? where id = ?;
|
|
)"_s;
|
|
|
|
static const QString DELETE_DOCUMENTS_SQL = uR"(
|
|
delete from documents where id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_DOCUMENT_SQL = uR"(
|
|
select id, document_time from documents where document_path = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_DOCUMENTS_SQL = uR"(
|
|
select id from documents where folder_id = ?;
|
|
)"_s;
|
|
|
|
static const QString SELECT_ALL_DOCUMENTS_SQL = uR"(
|
|
select id, document_path from documents;
|
|
)"_s;
|
|
|
|
static const QString SELECT_COUNT_STATISTICS_SQL = uR"(
|
|
select count(distinct d.id), sum(c.words), sum(c.tokens)
|
|
from documents d
|
|
left join chunks c on d.id = c.document_id
|
|
where d.folder_id = ?;
|
|
)"_s;
|
|
|
|
static bool addDocument(QSqlQuery &q, int folder_id, qint64 document_time, const QString &document_path, int *document_id)
|
|
{
|
|
if (!q.prepare(INSERT_DOCUMENTS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_id);
|
|
q.addBindValue(document_time);
|
|
q.addBindValue(document_path);
|
|
if (!q.exec())
|
|
return false;
|
|
*document_id = q.lastInsertId().toInt();
|
|
return true;
|
|
}
|
|
|
|
static bool removeDocument(QSqlQuery &q, int document_id)
|
|
{
|
|
if (!q.prepare(DELETE_DOCUMENTS_SQL))
|
|
return false;
|
|
q.addBindValue(document_id);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool updateDocument(QSqlQuery &q, int id, qint64 document_time)
|
|
{
|
|
if (!q.prepare(UPDATE_DOCUMENT_TIME_SQL))
|
|
return false;
|
|
q.addBindValue(document_time);
|
|
q.addBindValue(id);
|
|
return q.exec();
|
|
}
|
|
|
|
static bool selectDocument(QSqlQuery &q, const QString &document_path, int *id, qint64 *document_time)
|
|
{
|
|
if (!q.prepare(SELECT_DOCUMENT_SQL))
|
|
return false;
|
|
q.addBindValue(document_path);
|
|
if (!q.exec())
|
|
return false;
|
|
Q_ASSERT(q.size() < 2);
|
|
if (q.next()) {
|
|
*id = q.value(0).toInt();
|
|
*document_time = q.value(1).toLongLong();
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static bool selectDocuments(QSqlQuery &q, int folder_id, QList<int> *documentIds)
|
|
{
|
|
if (!q.prepare(SELECT_DOCUMENTS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_id);
|
|
if (!q.exec())
|
|
return false;
|
|
while (q.next())
|
|
documentIds->append(q.value(0).toInt());
|
|
return true;
|
|
}
|
|
|
|
static bool selectCountStatistics(QSqlQuery &q, int folder_id, int *total_docs, int *total_words, int *total_tokens)
|
|
{
|
|
if (!q.prepare(SELECT_COUNT_STATISTICS_SQL))
|
|
return false;
|
|
q.addBindValue(folder_id);
|
|
if (!q.exec())
|
|
return false;
|
|
if (q.next()) {
|
|
*total_docs = q.value(0).toInt();
|
|
*total_words = q.value(1).toInt();
|
|
*total_tokens = q.value(2).toInt();
|
|
}
|
|
return true;
|
|
}
|
|
|
|
// insert embedding only if still needed
|
|
static const QString INSERT_EMBEDDING_SQL = uR"(
|
|
insert into embeddings(model, folder_id, chunk_id, embedding)
|
|
select :model, d.folder_id, :chunk_id, :embedding
|
|
from chunks c
|
|
join documents d on d.id = c.document_id
|
|
join folders f on f.id = d.folder_id
|
|
join collection_items ci on ci.folder_id = f.id
|
|
join collections co on co.id = ci.collection_id
|
|
where co.embedding_model = :model and c.id = :chunk_id
|
|
limit 1;
|
|
)"_s;
|
|
|
|
static const QString GET_COLLECTION_EMBEDDINGS_SQL = uR"(
|
|
select e.chunk_id, e.embedding
|
|
from embeddings e
|
|
join collections co on co.embedding_model = e.model
|
|
join collection_items ci on ci.folder_id = e.folder_id and ci.collection_id = co.id
|
|
where co.name in ('%1');
|
|
)"_s;
|
|
|
|
static const QString GET_CHUNK_FILE_SQL = uR"(
|
|
select file from chunks where id = ?;
|
|
)"_s;
|
|
|
|
namespace {
|
|
struct Embedding { QString model; int folder_id; int chunk_id; QByteArray data; };
|
|
struct EmbeddingStat { QString lastFile; int nAdded; int nSkipped; };
|
|
} // namespace
|
|
|
|
NAMED_PAIR(EmbeddingFolder, QString, embedding_model, int, folder_id)
|
|
|
|
static bool sqlAddEmbeddings(QSqlQuery &q, const QList<Embedding> &embeddings, QHash<EmbeddingFolder, EmbeddingStat> &embeddingStats)
|
|
{
|
|
if (!q.prepare(INSERT_EMBEDDING_SQL))
|
|
return false;
|
|
|
|
// insert embedding if needed
|
|
for (const auto &e: embeddings) {
|
|
q.bindValue(":model", e.model);
|
|
q.bindValue(":chunk_id", e.chunk_id);
|
|
q.bindValue(":embedding", e.data);
|
|
if (!q.exec())
|
|
return false;
|
|
|
|
auto &stat = embeddingStats[{ e.model, e.folder_id }];
|
|
if (q.numRowsAffected()) {
|
|
stat.nAdded++; // embedding added
|
|
} else {
|
|
stat.nSkipped++; // embedding no longer needed
|
|
}
|
|
}
|
|
|
|
if (!q.prepare(GET_CHUNK_FILE_SQL))
|
|
return false;
|
|
|
|
// populate statistics for each collection item
|
|
for (const auto &e: embeddings) {
|
|
auto &stat = embeddingStats[{ e.model, e.folder_id }];
|
|
if (stat.nAdded && stat.lastFile.isNull()) {
|
|
q.addBindValue(e.chunk_id);
|
|
if (!q.exec() || !q.next())
|
|
return false;
|
|
stat.lastFile = q.value(0).toString();
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
void Database::transaction()
|
|
{
|
|
bool ok = m_db.transaction();
|
|
Q_ASSERT(ok);
|
|
}
|
|
|
|
void Database::commit()
|
|
{
|
|
bool ok = m_db.commit();
|
|
Q_ASSERT(ok);
|
|
}
|
|
|
|
void Database::rollback()
|
|
{
|
|
bool ok = m_db.rollback();
|
|
Q_ASSERT(ok);
|
|
}
|
|
|
|
bool Database::hasContent()
|
|
{
|
|
return m_db.tables().contains("chunks", Qt::CaseInsensitive);
|
|
}
|
|
|
|
int Database::openDatabase(const QString &modelPath, bool create, int ver)
|
|
{
|
|
if (!QFileInfo(modelPath).isDir()) {
|
|
qWarning() << "ERROR: invalid download path" << modelPath;
|
|
return -1;
|
|
}
|
|
if (m_db.isOpen())
|
|
m_db.close();
|
|
auto dbPath = u"%1/localdocs_v%2.db"_s.arg(modelPath).arg(ver);
|
|
if (!create && !QFileInfo::exists(dbPath))
|
|
return 0;
|
|
m_db.setDatabaseName(dbPath);
|
|
if (!m_db.open()) {
|
|
qWarning() << "ERROR: opening db" << dbPath << m_db.lastError();
|
|
return -1;
|
|
}
|
|
return hasContent();
|
|
}
|
|
|
|
bool Database::openLatestDb(const QString &modelPath, QList<CollectionItem> &oldCollections)
|
|
{
|
|
/*
|
|
* Support upgrade path from older versions:
|
|
*
|
|
* 1. Detect and load dbPath with older versions
|
|
* 2. Provide versioned SQL select statements
|
|
* 3. Upgrade the tables to the new version
|
|
* 4. By default mark all collections of older versions as force indexing and present to the user
|
|
* the an 'update' button letting them know a breaking change happened and that the collection
|
|
* will need to be indexed again
|
|
*/
|
|
|
|
int dbVer;
|
|
for (dbVer = LOCALDOCS_VERSION;; dbVer--) {
|
|
if (dbVer < LOCALDOCS_MIN_VER) return true; // create a new db
|
|
int res = openDatabase(modelPath, false, dbVer);
|
|
if (res == 1) break; // found one with content
|
|
if (res == -1) return false; // error
|
|
}
|
|
|
|
if (dbVer == LOCALDOCS_VERSION) return true; // already up-to-date
|
|
|
|
// If we're upgrading, then we need to do a select on the current version of the collections table,
|
|
// then create the new one and populate the collections table and mark them as needing forced
|
|
// indexing
|
|
|
|
#if defined(DEBUG)
|
|
qDebug() << "Older localdocs version found" << dbVer << "upgrade to" << LOCALDOCS_VERSION;
|
|
#endif
|
|
|
|
// Select the current collections which will be marked to force indexing
|
|
QSqlQuery q(m_db);
|
|
if (!selectAllFromCollections(q, &oldCollections, dbVer)) {
|
|
qWarning() << "ERROR: Could not open select old collections" << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
m_db.close();
|
|
return true;
|
|
}
|
|
|
|
bool Database::initDb(const QString &modelPath, const QList<CollectionItem> &oldCollections)
|
|
{
|
|
if (!m_db.isOpen()) {
|
|
int res = openDatabase(modelPath);
|
|
if (res == 1) return true; // already populated
|
|
if (res == -1) return false; // error
|
|
} else if (hasContent()) {
|
|
return true; // already populated
|
|
}
|
|
|
|
transaction();
|
|
|
|
QSqlQuery q(m_db);
|
|
for (const auto &cmd: INIT_DB_SQL) {
|
|
if (!q.exec(cmd)) {
|
|
qWarning() << "ERROR: failed to create tables" << q.lastError();
|
|
rollback();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/* These are collection items that came from an older version of localdocs which
|
|
* require forced indexing that should only be done when the user has explicitly asked
|
|
* for them to be indexed again */
|
|
for (const CollectionItem &item : oldCollections) {
|
|
if (!addFolder(item.collection, item.folder_path, QString())) {
|
|
qWarning() << "ERROR: failed to add previous collections to new database";
|
|
rollback();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
commit();
|
|
return true;
|
|
}
|
|
|
|
Database::Database(int chunkSize, QStringList extensions)
|
|
: QObject(nullptr)
|
|
, m_chunkSize(chunkSize)
|
|
, m_scannedFileExtensions(std::move(extensions))
|
|
, m_scanTimer(new QTimer(this))
|
|
, m_watcher(new QFileSystemWatcher(this))
|
|
, m_embLLM(new EmbeddingLLM)
|
|
, m_databaseValid(true)
|
|
{
|
|
m_db = QSqlDatabase::database(QSqlDatabase::defaultConnection, false);
|
|
if (!m_db.isValid())
|
|
m_db = QSqlDatabase::addDatabase("QSQLITE");
|
|
Q_ASSERT(m_db.isValid());
|
|
|
|
moveToThread(&m_dbThread);
|
|
m_dbThread.setObjectName("database");
|
|
m_dbThread.start();
|
|
}
|
|
|
|
Database::~Database()
|
|
{
|
|
m_dbThread.quit();
|
|
m_dbThread.wait();
|
|
delete m_embLLM;
|
|
}
|
|
|
|
void Database::setStartUpdateTime(CollectionItem &item)
|
|
{
|
|
QSqlQuery q(m_db);
|
|
const qint64 update_time = QDateTime::currentMSecsSinceEpoch();
|
|
if (!updateStartUpdateTime(q, item.collection_id, update_time))
|
|
qWarning() << "Database ERROR: failed to set start update time:" << q.lastError();
|
|
else
|
|
item.startUpdate = QDateTime::fromMSecsSinceEpoch(update_time);
|
|
}
|
|
|
|
void Database::setLastUpdateTime(CollectionItem &item)
|
|
{
|
|
QSqlQuery q(m_db);
|
|
const qint64 update_time = QDateTime::currentMSecsSinceEpoch();
|
|
if (!updateLastUpdateTime(q, item.collection_id, update_time))
|
|
qWarning() << "Database ERROR: failed to set last update time:" << q.lastError();
|
|
else
|
|
item.lastUpdate = QDateTime::fromMSecsSinceEpoch(update_time);
|
|
}
|
|
|
|
CollectionItem Database::guiCollectionItem(int folder_id) const
|
|
{
|
|
Q_ASSERT(m_collectionMap.contains(folder_id));
|
|
return m_collectionMap.value(folder_id);
|
|
}
|
|
|
|
void Database::updateGuiForCollectionItem(const CollectionItem &item)
|
|
{
|
|
m_collectionMap.insert(item.folder_id, item);
|
|
emit requestUpdateGuiForCollectionItem(item);
|
|
}
|
|
|
|
void Database::addGuiCollectionItem(const CollectionItem &item)
|
|
{
|
|
m_collectionMap.insert(item.folder_id, item);
|
|
emit requestAddGuiCollectionItem(item);
|
|
}
|
|
|
|
void Database::removeGuiFolderById(const QString &collection, int folder_id)
|
|
{
|
|
emit requestRemoveGuiFolderById(collection, folder_id);
|
|
}
|
|
|
|
void Database::guiCollectionListUpdated(const QList<CollectionItem> &collectionList)
|
|
{
|
|
for (const auto &i : collectionList)
|
|
m_collectionMap.insert(i.folder_id, i);
|
|
emit requestGuiCollectionListUpdated(collectionList);
|
|
}
|
|
|
|
void Database::updateFolderToIndex(int folder_id, size_t countForFolder, bool sendChunks)
|
|
{
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.currentDocsToIndex = countForFolder;
|
|
if (!countForFolder) {
|
|
if (sendChunks && !m_chunkList.isEmpty())
|
|
sendChunkList(); // send any remaining embedding chunks to llm
|
|
item.indexing = false;
|
|
item.installed = true;
|
|
|
|
// Set the last update if we are done
|
|
if (item.startUpdate > item.lastUpdate && item.currentEmbeddingsToIndex == 0)
|
|
setLastUpdateTime(item);
|
|
}
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
|
|
void Database::handleDocumentError(const QString &errorMessage,
|
|
int document_id, const QString &document_path, const QSqlError &error)
|
|
{
|
|
qWarning() << errorMessage << document_id << document_path << error;
|
|
}
|
|
|
|
size_t Database::chunkStream(QTextStream &stream, int folder_id, int document_id, const QString &embedding_model,
|
|
const QString &file, const QString &title, const QString &author, const QString &subject, const QString &keywords,
|
|
int page, int maxChunks)
|
|
{
|
|
int charCount = 0;
|
|
// TODO: implement line_from/line_to
|
|
constexpr int line_from = -1;
|
|
constexpr int line_to = -1;
|
|
QList<QString> words;
|
|
int chunks = 0;
|
|
int addedWords = 0;
|
|
|
|
for (;;) {
|
|
QString word;
|
|
stream >> word;
|
|
if (stream.status() && !stream.atEnd())
|
|
return -1;
|
|
charCount += word.length();
|
|
if (!word.isEmpty())
|
|
words.append(word);
|
|
if (stream.status() || charCount + words.size() - 1 >= m_chunkSize) {
|
|
if (!words.isEmpty()) {
|
|
const QString chunk = words.join(" ");
|
|
QSqlQuery q(m_db);
|
|
int chunk_id = 0;
|
|
if (!addChunk(q,
|
|
document_id,
|
|
chunk,
|
|
file,
|
|
title,
|
|
author,
|
|
subject,
|
|
keywords,
|
|
page,
|
|
line_from,
|
|
line_to,
|
|
words.size(),
|
|
&chunk_id
|
|
)) {
|
|
qWarning() << "ERROR: Could not insert chunk into db" << q.lastError();
|
|
}
|
|
|
|
addedWords += words.size();
|
|
|
|
EmbeddingChunk toEmbed;
|
|
toEmbed.model = embedding_model;
|
|
toEmbed.folder_id = folder_id;
|
|
toEmbed.chunk_id = chunk_id;
|
|
toEmbed.chunk = chunk;
|
|
appendChunk(toEmbed);
|
|
++chunks;
|
|
|
|
words.clear();
|
|
charCount = 0;
|
|
}
|
|
|
|
if (stream.status() || (maxChunks > 0 && chunks == maxChunks))
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (chunks) {
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
|
|
// Set the start update if we haven't done so already
|
|
if (item.startUpdate <= item.lastUpdate && item.currentEmbeddingsToIndex == 0)
|
|
setStartUpdateTime(item);
|
|
|
|
item.currentEmbeddingsToIndex += chunks;
|
|
item.totalEmbeddingsToIndex += chunks;
|
|
item.totalWords += addedWords;
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
|
|
return stream.pos();
|
|
}
|
|
|
|
void Database::appendChunk(const EmbeddingChunk &chunk)
|
|
{
|
|
m_chunkList.reserve(s_batchSize);
|
|
m_chunkList.append(chunk);
|
|
if (m_chunkList.size() >= s_batchSize)
|
|
sendChunkList();
|
|
}
|
|
|
|
void Database::sendChunkList()
|
|
{
|
|
m_embLLM->generateDocEmbeddingsAsync(m_chunkList);
|
|
m_chunkList.clear();
|
|
}
|
|
|
|
void Database::handleEmbeddingsGenerated(const QVector<EmbeddingResult> &embeddings)
|
|
{
|
|
Q_ASSERT(!embeddings.isEmpty());
|
|
|
|
QList<Embedding> sqlEmbeddings;
|
|
for (const auto &e: embeddings) {
|
|
auto data = QByteArray::fromRawData(
|
|
reinterpret_cast<const char *>(e.embedding.data()),
|
|
e.embedding.size() * sizeof(e.embedding.front())
|
|
);
|
|
sqlEmbeddings.append({e.model, e.folder_id, e.chunk_id, std::move(data)});
|
|
}
|
|
|
|
transaction();
|
|
|
|
QSqlQuery q(m_db);
|
|
QHash<EmbeddingFolder, EmbeddingStat> stats;
|
|
if (!sqlAddEmbeddings(q, sqlEmbeddings, stats)) {
|
|
qWarning() << "Database ERROR: failed to add embeddings:" << q.lastError();
|
|
return rollback();
|
|
}
|
|
|
|
commit();
|
|
|
|
// FIXME(jared): embedding counts are per-collectionitem, not per-folder
|
|
for (const auto &[key, stat]: std::as_const(stats).asKeyValueRange()) {
|
|
if (!m_collectionMap.contains(key.folder_id)) continue;
|
|
CollectionItem item = guiCollectionItem(key.folder_id);
|
|
item.currentEmbeddingsToIndex -= stat.nAdded + stat.nSkipped;
|
|
item.totalEmbeddingsToIndex -= stat.nSkipped;
|
|
if (!stat.lastFile.isNull())
|
|
item.fileCurrentlyProcessing = stat.lastFile;
|
|
|
|
// Set the last update if we are done
|
|
Q_ASSERT(item.startUpdate > item.lastUpdate);
|
|
if (!item.indexing && item.currentEmbeddingsToIndex == 0)
|
|
setLastUpdateTime(item);
|
|
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
}
|
|
|
|
void Database::handleErrorGenerated(const QVector<EmbeddingChunk> &chunks, const QString &error)
|
|
{
|
|
/* FIXME(jared): errors are actually collection-specific because they are conditioned
|
|
* on the embedding model, but this sets the error on all collections for a given
|
|
* folder */
|
|
|
|
QSet<int> folder_ids;
|
|
for (const auto &c: chunks) { folder_ids << c.folder_id; }
|
|
|
|
for (int fid: folder_ids) {
|
|
if (!m_collectionMap.contains(fid)) continue;
|
|
CollectionItem item = guiCollectionItem(fid);
|
|
item.error = error;
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
}
|
|
|
|
size_t Database::countOfDocuments(int folder_id) const
|
|
{
|
|
if (!m_docsToScan.contains(folder_id))
|
|
return 0;
|
|
return m_docsToScan.value(folder_id).size();
|
|
}
|
|
|
|
size_t Database::countOfBytes(int folder_id) const
|
|
{
|
|
if (!m_docsToScan.contains(folder_id))
|
|
return 0;
|
|
size_t totalBytes = 0;
|
|
const QQueue<DocumentInfo> &docs = m_docsToScan.value(folder_id);
|
|
for (const DocumentInfo &f : docs)
|
|
totalBytes += f.doc.size();
|
|
return totalBytes;
|
|
}
|
|
|
|
DocumentInfo Database::dequeueDocument()
|
|
{
|
|
Q_ASSERT(!m_docsToScan.isEmpty());
|
|
const int firstKey = m_docsToScan.firstKey();
|
|
QQueue<DocumentInfo> &queue = m_docsToScan[firstKey];
|
|
Q_ASSERT(!queue.isEmpty());
|
|
DocumentInfo result = queue.dequeue();
|
|
if (queue.isEmpty())
|
|
m_docsToScan.remove(firstKey);
|
|
return result;
|
|
}
|
|
|
|
void Database::removeFolderFromDocumentQueue(int folder_id)
|
|
{
|
|
if (!m_docsToScan.contains(folder_id))
|
|
return;
|
|
m_docsToScan.remove(folder_id);
|
|
}
|
|
|
|
void Database::enqueueDocumentInternal(const DocumentInfo &info, bool prepend)
|
|
{
|
|
const int key = info.folder;
|
|
if (!m_docsToScan.contains(key))
|
|
m_docsToScan[key] = QQueue<DocumentInfo>();
|
|
if (prepend)
|
|
m_docsToScan[key].prepend(info);
|
|
else
|
|
m_docsToScan[key].enqueue(info);
|
|
}
|
|
|
|
void Database::enqueueDocuments(int folder_id, const QVector<DocumentInfo> &infos)
|
|
{
|
|
for (int i = 0; i < infos.size(); ++i)
|
|
enqueueDocumentInternal(infos[i]);
|
|
const size_t count = countOfDocuments(folder_id);
|
|
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.currentDocsToIndex = count;
|
|
item.totalDocsToIndex = count;
|
|
const size_t bytes = countOfBytes(folder_id);
|
|
item.currentBytesToIndex = bytes;
|
|
item.totalBytesToIndex = bytes;
|
|
updateGuiForCollectionItem(item);
|
|
m_scanTimer->start();
|
|
}
|
|
|
|
void Database::scanQueueBatch()
|
|
{
|
|
QElapsedTimer timer;
|
|
timer.start();
|
|
|
|
transaction();
|
|
|
|
// scan for up to 100ms or until we run out of documents
|
|
while (!m_docsToScan.isEmpty() && timer.elapsed() < 100)
|
|
scanQueue();
|
|
|
|
commit();
|
|
|
|
if (m_docsToScan.isEmpty())
|
|
m_scanTimer->stop();
|
|
}
|
|
|
|
void Database::scanQueue()
|
|
{
|
|
DocumentInfo info = dequeueDocument();
|
|
const size_t countForFolder = countOfDocuments(info.folder);
|
|
const int folder_id = info.folder;
|
|
|
|
// Update info
|
|
info.doc.stat();
|
|
|
|
// If the doc has since been deleted or no longer readable, then we schedule more work and return
|
|
// leaving the cleanup for the cleanup handler
|
|
if (!info.doc.exists() || !info.doc.isReadable())
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
|
|
const qint64 document_time = info.doc.fileTime(QFile::FileModificationTime).toMSecsSinceEpoch();
|
|
const QString document_path = info.doc.canonicalFilePath();
|
|
const bool currentlyProcessing = info.currentlyProcessing;
|
|
|
|
// Check and see if we already have this document
|
|
QSqlQuery q(m_db);
|
|
int existing_id = -1;
|
|
qint64 existing_time = -1;
|
|
if (!selectDocument(q, document_path, &existing_id, &existing_time)) {
|
|
handleDocumentError("ERROR: Cannot select document",
|
|
existing_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
|
|
// If we have the document, we need to compare the last modification time and if it is newer
|
|
// we must rescan the document, otherwise return
|
|
if (existing_id != -1 && !currentlyProcessing) {
|
|
Q_ASSERT(existing_time != -1);
|
|
if (document_time == existing_time) {
|
|
// No need to rescan, but we do have to schedule next
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
if (!removeChunksByDocumentId(q, existing_id)) {
|
|
handleDocumentError("ERROR: Cannot remove chunks of document",
|
|
existing_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
updateCollectionStatistics();
|
|
}
|
|
|
|
// Update the document_time for an existing document, or add it for the first time now
|
|
int document_id = existing_id;
|
|
if (!currentlyProcessing) {
|
|
if (document_id != -1) {
|
|
if (!updateDocument(q, document_id, document_time)) {
|
|
handleDocumentError("ERROR: Could not update document_time",
|
|
document_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
} else {
|
|
if (!addDocument(q, folder_id, document_time, document_path, &document_id)) {
|
|
handleDocumentError("ERROR: Could not add document",
|
|
document_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.totalDocs += 1;
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
}
|
|
|
|
// Get the embedding model for this folder
|
|
// FIXME(jared): there can be more than one since we allow one folder to be in multiple collections
|
|
QString embedding_model;
|
|
if (!sqlGetFolderEmbeddingModel(q, folder_id, embedding_model)) {
|
|
handleDocumentError("ERROR: Could not get embedding model",
|
|
document_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
|
|
Q_ASSERT(document_id != -1);
|
|
if (info.isPdf()) {
|
|
QPdfDocument doc;
|
|
if (QPdfDocument::Error::None != doc.load(info.doc.canonicalFilePath())) {
|
|
handleDocumentError("ERROR: Could not load pdf",
|
|
document_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
const size_t bytes = info.doc.size();
|
|
const size_t bytesPerPage = std::floor(bytes / doc.pageCount());
|
|
const int pageIndex = info.currentPage;
|
|
#if defined(DEBUG)
|
|
qDebug() << "scanning page" << pageIndex << "of" << doc.pageCount() << document_path;
|
|
#endif
|
|
const QPdfSelection selection = doc.getAllText(pageIndex);
|
|
QString text = selection.text();
|
|
QTextStream stream(&text);
|
|
chunkStream(stream, info.folder, document_id, embedding_model, info.doc.fileName(),
|
|
doc.metaData(QPdfDocument::MetaDataField::Title).toString(),
|
|
doc.metaData(QPdfDocument::MetaDataField::Author).toString(),
|
|
doc.metaData(QPdfDocument::MetaDataField::Subject).toString(),
|
|
doc.metaData(QPdfDocument::MetaDataField::Keywords).toString(),
|
|
pageIndex + 1
|
|
);
|
|
CollectionItem item = guiCollectionItem(info.folder);
|
|
item.currentBytesToIndex -= bytesPerPage;
|
|
updateGuiForCollectionItem(item);
|
|
if (info.currentPage < doc.pageCount()) {
|
|
info.currentPage += 1;
|
|
info.currentlyProcessing = true;
|
|
enqueueDocumentInternal(info, true /*prepend*/);
|
|
return updateFolderToIndex(folder_id, countForFolder + 1);
|
|
}
|
|
|
|
item.currentBytesToIndex -= bytes - (bytesPerPage * doc.pageCount());
|
|
updateGuiForCollectionItem(item);
|
|
} else {
|
|
BinaryDetectingFile file(document_path);
|
|
if (!file.open(QIODevice::ReadOnly)) {
|
|
handleDocumentError("ERROR: Cannot open file for scanning",
|
|
existing_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
Q_ASSERT(!file.isSequential()); // we need to seek
|
|
|
|
const size_t bytes = info.doc.size();
|
|
QTextStream stream(&file);
|
|
const size_t byteIndex = info.currentPosition;
|
|
if (byteIndex) {
|
|
/* Read the Unicode BOM to detect the encoding. Without this, QTextStream will
|
|
* always interpret the text as UTF-8 when byteIndex is nonzero. */
|
|
stream.read(1);
|
|
|
|
if (!stream.seek(byteIndex)) {
|
|
handleDocumentError("ERROR: Cannot seek to pos for scanning",
|
|
existing_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
}
|
|
#if defined(DEBUG)
|
|
qDebug() << "scanning byteIndex" << byteIndex << "of" << bytes << document_path;
|
|
#endif
|
|
int pos = chunkStream(stream, info.folder, document_id, embedding_model, info.doc.fileName(),
|
|
QString() /*title*/, QString() /*author*/, QString() /*subject*/, QString() /*keywords*/, -1 /*page*/,
|
|
100 /*maxChunks*/);
|
|
if (pos < 0) {
|
|
if (!file.binarySeen()) {
|
|
handleDocumentError(u"ERROR: Failed to read file (status %1)"_s.arg(stream.status()),
|
|
existing_id, document_path, q.lastError());
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
|
|
/* When we see a binary file, we treat it like an empty file so we know not to
|
|
* scan it again. All existing chunks are removed, and in-progress embeddings
|
|
* are ignored when they complete. */
|
|
|
|
qInfo() << "LocalDocs: Ignoring file with binary data:" << document_path;
|
|
|
|
// this will also ensure in-flight embeddings are ignored
|
|
if (!removeChunksByDocumentId(q, existing_id)) {
|
|
handleDocumentError("ERROR: Cannot remove chunks of document",
|
|
existing_id, document_path, q.lastError());
|
|
}
|
|
updateCollectionStatistics();
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
file.close();
|
|
const size_t bytesChunked = pos - byteIndex;
|
|
CollectionItem item = guiCollectionItem(info.folder);
|
|
item.currentBytesToIndex -= bytesChunked;
|
|
updateGuiForCollectionItem(item);
|
|
if (info.currentPosition < bytes) {
|
|
info.currentPosition = pos;
|
|
info.currentlyProcessing = true;
|
|
enqueueDocumentInternal(info, true /*prepend*/);
|
|
return updateFolderToIndex(folder_id, countForFolder + 1);
|
|
}
|
|
}
|
|
|
|
return updateFolderToIndex(folder_id, countForFolder);
|
|
}
|
|
|
|
void Database::scanDocuments(int folder_id, const QString &folder_path)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "scanning folder for documents" << folder_path;
|
|
#endif
|
|
|
|
QDirIterator it(folder_path, QDir::Readable | QDir::Files | QDir::Dirs | QDir::NoDotAndDotDot,
|
|
QDirIterator::Subdirectories);
|
|
QVector<DocumentInfo> infos;
|
|
while (it.hasNext()) {
|
|
it.next();
|
|
QFileInfo fileInfo = it.fileInfo();
|
|
if (fileInfo.isDir()) {
|
|
addFolderToWatch(fileInfo.canonicalFilePath());
|
|
continue;
|
|
}
|
|
|
|
if (!m_scannedFileExtensions.contains(fileInfo.suffix(), Qt::CaseInsensitive))
|
|
continue;
|
|
|
|
DocumentInfo info;
|
|
info.folder = folder_id;
|
|
info.doc = fileInfo;
|
|
infos.append(info);
|
|
}
|
|
|
|
if (!infos.isEmpty()) {
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.indexing = true;
|
|
updateGuiForCollectionItem(item);
|
|
enqueueDocuments(folder_id, infos);
|
|
} else {
|
|
updateFolderToIndex(folder_id, 0, false);
|
|
}
|
|
}
|
|
|
|
void Database::start()
|
|
{
|
|
connect(m_watcher, &QFileSystemWatcher::directoryChanged, this, &Database::directoryChanged);
|
|
connect(m_embLLM, &EmbeddingLLM::embeddingsGenerated, this, &Database::handleEmbeddingsGenerated);
|
|
connect(m_embLLM, &EmbeddingLLM::errorGenerated, this, &Database::handleErrorGenerated);
|
|
m_scanTimer->callOnTimeout(this, &Database::scanQueueBatch);
|
|
|
|
const QString modelPath = MySettings::globalInstance()->modelPath();
|
|
QList<CollectionItem> oldCollections;
|
|
|
|
if (!openLatestDb(modelPath, oldCollections)) {
|
|
m_databaseValid = false;
|
|
} else if (!initDb(modelPath, oldCollections)) {
|
|
m_databaseValid = false;
|
|
} else {
|
|
cleanDB();
|
|
addCurrentFolders();
|
|
}
|
|
|
|
if (!m_databaseValid)
|
|
emit databaseValidChanged();
|
|
}
|
|
|
|
void Database::addCurrentFolders()
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "addCurrentFolders";
|
|
#endif
|
|
|
|
QSqlQuery q(m_db);
|
|
QList<CollectionItem> collections;
|
|
if (!selectAllFromCollections(q, &collections)) {
|
|
qWarning() << "ERROR: Cannot select collections" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
guiCollectionListUpdated(collections);
|
|
|
|
scheduleUncompletedEmbeddings();
|
|
|
|
for (const auto &i : collections) {
|
|
if (!i.forceIndexing) {
|
|
addFolderToWatch(i.folder_path);
|
|
scanDocuments(i.folder_id, i.folder_path);
|
|
}
|
|
}
|
|
|
|
updateCollectionStatistics();
|
|
}
|
|
|
|
void Database::scheduleUncompletedEmbeddings()
|
|
{
|
|
QHash<IncompleteChunk, QStringList> chunkList;
|
|
QSqlQuery q(m_db);
|
|
if (!selectAllUncompletedChunks(q, chunkList)) {
|
|
qWarning() << "ERROR: Cannot select uncompleted chunks" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
if (chunkList.isEmpty())
|
|
return;
|
|
|
|
// map of folder_id -> chunk count
|
|
QMap<int, int> folderNChunks;
|
|
for (auto it = chunkList.keyBegin(), end = chunkList.keyEnd(); it != end; ++it) {
|
|
int folder_id = it->folder_id;
|
|
|
|
if (folderNChunks.contains(folder_id)) continue;
|
|
int total = 0;
|
|
if (!selectCountChunks(q, folder_id, total)) {
|
|
qWarning() << "ERROR: Cannot count total chunks" << q.lastError();
|
|
return;
|
|
}
|
|
folderNChunks.insert(folder_id, total);
|
|
}
|
|
|
|
// map of (folder_id, collection) -> incomplete count
|
|
QMap<QPair<int, QString>, int> itemNIncomplete;
|
|
for (const auto &[chunk, collections]: std::as_const(chunkList).asKeyValueRange())
|
|
for (const auto &collection: std::as_const(collections))
|
|
itemNIncomplete[{ chunk.folder_id, collection }]++;
|
|
|
|
for (const auto &[key, nIncomplete]: std::as_const(itemNIncomplete).asKeyValueRange()) {
|
|
const auto &[folder_id, collection] = key;
|
|
|
|
/* FIXME(jared): this needs to be split by collection because different
|
|
* collections have different embedding models */
|
|
int total = folderNChunks.value(folder_id);
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.totalEmbeddingsToIndex = total;
|
|
item.currentEmbeddingsToIndex = nIncomplete;
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
|
|
for (auto it = chunkList.keyBegin(), end = chunkList.keyEnd(); it != end;) {
|
|
QList<EmbeddingChunk> batch;
|
|
for (; it != end && batch.size() < s_batchSize; ++it)
|
|
batch.append({ /*model*/ it->embedding_model, /*folder_id*/ it->folder_id, /*chunk_id*/ it->chunk_id, /*chunk*/ it->text });
|
|
Q_ASSERT(!batch.isEmpty());
|
|
m_embLLM->generateDocEmbeddingsAsync(batch);
|
|
}
|
|
}
|
|
|
|
void Database::updateCollectionStatistics()
|
|
{
|
|
QSqlQuery q(m_db);
|
|
QList<CollectionItem> collections;
|
|
if (!selectAllFromCollections(q, &collections)) {
|
|
qWarning() << "ERROR: Cannot select collections" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
for (const auto &i: std::as_const(collections)) {
|
|
int total_docs = 0;
|
|
int total_words = 0;
|
|
int total_tokens = 0;
|
|
if (!selectCountStatistics(q, i.folder_id, &total_docs, &total_words, &total_tokens)) {
|
|
qWarning() << "ERROR: could not count statistics for folder" << q.lastError();
|
|
} else {
|
|
CollectionItem item = guiCollectionItem(i.folder_id);
|
|
item.totalDocs = total_docs;
|
|
item.totalWords = total_words;
|
|
item.totalTokens = total_tokens;
|
|
updateGuiForCollectionItem(item);
|
|
}
|
|
}
|
|
}
|
|
|
|
int Database::checkAndAddFolderToDB(const QString &path)
|
|
{
|
|
QFileInfo info(path);
|
|
if (!info.exists() || !info.isReadable()) {
|
|
qWarning() << "ERROR: Cannot add folder that doesn't exist or not readable" << path;
|
|
return -1;
|
|
}
|
|
|
|
QSqlQuery q(m_db);
|
|
int folder_id = -1;
|
|
|
|
// See if the folder exists in the db
|
|
if (!selectFolder(q, path, &folder_id)) {
|
|
qWarning() << "ERROR: Cannot select folder from path" << path << q.lastError();
|
|
return -1;
|
|
}
|
|
|
|
// Add the folder
|
|
if (folder_id == -1 && !addFolderToDB(q, path, &folder_id)) {
|
|
qWarning() << "ERROR: Cannot add folder to db with path" << path << q.lastError();
|
|
return -1;
|
|
}
|
|
|
|
Q_ASSERT(folder_id != -1);
|
|
return folder_id;
|
|
}
|
|
|
|
void Database::forceIndexing(const QString &collection, const QString &embedding_model)
|
|
{
|
|
Q_ASSERT(!embedding_model.isNull());
|
|
|
|
QSqlQuery q(m_db);
|
|
QList<QPair<int, QString>> folders;
|
|
if (!selectFoldersFromCollection(q, collection, &folders)) {
|
|
qWarning() << "ERROR: Cannot select folders from collections" << collection << q.lastError();
|
|
return;
|
|
}
|
|
|
|
if (!setCollectionEmbeddingModel(q, collection, embedding_model)) {
|
|
qWarning().nospace() << "ERROR: Cannot set embedding model for collection " << collection << ": "
|
|
<< q.lastError();
|
|
return;
|
|
}
|
|
|
|
for (const auto &folder: std::as_const(folders)) {
|
|
CollectionItem item = guiCollectionItem(folder.first);
|
|
item.embeddingModel = embedding_model;
|
|
item.forceIndexing = false;
|
|
updateGuiForCollectionItem(item);
|
|
addFolderToWatch(folder.second);
|
|
scanDocuments(folder.first, folder.second);
|
|
}
|
|
}
|
|
|
|
void Database::forceRebuildFolder(const QString &path)
|
|
{
|
|
QSqlQuery q(m_db);
|
|
int folder_id;
|
|
if (!selectFolder(q, path, &folder_id)) {
|
|
qWarning().nospace() << "Database ERROR: Cannot select folder from path " << path << ": " << q.lastError();
|
|
return;
|
|
}
|
|
|
|
Q_ASSERT(!m_docsToScan.contains(folder_id));
|
|
|
|
transaction();
|
|
|
|
if (!sqlRemoveDocsByFolderPath(q, path)) {
|
|
qWarning().nospace() << "Database ERROR: Cannot remove chunks for folder " << path << ": " << q.lastError();
|
|
return rollback();
|
|
}
|
|
|
|
commit();
|
|
|
|
updateCollectionStatistics();
|
|
|
|
// We now have zero embeddings. Document progress will be updated by scanDocuments.
|
|
// FIXME(jared): this updates the folder, but these values should also depend on the collection
|
|
CollectionItem item = guiCollectionItem(folder_id);
|
|
item.currentEmbeddingsToIndex = item.totalEmbeddingsToIndex = 0;
|
|
updateGuiForCollectionItem(item);
|
|
|
|
scanDocuments(folder_id, path);
|
|
}
|
|
|
|
bool Database::addFolder(const QString &collection, const QString &path, const QString &embedding_model)
|
|
{
|
|
// add the folder, if needed
|
|
const int folder_id = checkAndAddFolderToDB(path);
|
|
if (folder_id == -1)
|
|
return false;
|
|
|
|
std::optional<CollectionItem> item;
|
|
QSqlQuery q(m_db);
|
|
if (!selectCollectionByName(q, collection, item)) {
|
|
qWarning().nospace() << "Database ERROR: Cannot select collection " << collection << ": " << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
// add the collection, if needed
|
|
if (!item) {
|
|
item.emplace();
|
|
if (!addCollection(q, collection, QDateTime() /*start_update*/, QDateTime() /*last_update*/,
|
|
embedding_model /*embedding_model*/, *item)) {
|
|
qWarning().nospace() << "ERROR: Cannot add collection " << collection << ": " << q.lastError();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// link the folder and the collection, if needed
|
|
int res = addCollectionItem(q, item->collection_id, folder_id);
|
|
if (res < 0) { // error
|
|
qWarning().nospace() << "Database ERROR: Cannot add folder " << path << " to collection " << collection << ": "
|
|
<< q.lastError();
|
|
return false;
|
|
}
|
|
|
|
// add the new collection item to the UI
|
|
if (res == 1) { // new item added
|
|
item->folder_path = path;
|
|
item->folder_id = folder_id;
|
|
addGuiCollectionItem(item.value());
|
|
|
|
// note: this is the existing embedding model if the collection was found
|
|
if (!item->embeddingModel.isNull()) {
|
|
addFolderToWatch(path);
|
|
scanDocuments(folder_id, path);
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void Database::removeFolder(const QString &collection, const QString &path)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "removeFolder" << path;
|
|
#endif
|
|
|
|
QSqlQuery q(m_db);
|
|
int folder_id = -1;
|
|
|
|
// See if the folder exists in the db
|
|
if (!selectFolder(q, path, &folder_id)) {
|
|
qWarning() << "ERROR: Cannot select folder from path" << path << q.lastError();
|
|
return;
|
|
}
|
|
|
|
// If we don't have a folder_id in the db, then something bad has happened
|
|
Q_ASSERT(folder_id != -1);
|
|
if (folder_id == -1) {
|
|
qWarning() << "ERROR: Collected folder does not exist in db" << path;
|
|
m_watcher->removePath(path);
|
|
return;
|
|
}
|
|
|
|
transaction();
|
|
|
|
if (removeFolderInternal(collection, folder_id, path)) {
|
|
commit();
|
|
} else {
|
|
rollback();
|
|
}
|
|
}
|
|
|
|
bool Database::removeFolderInternal(const QString &collection, int folder_id, const QString &path)
|
|
{
|
|
// Remove it from the collection
|
|
QSqlQuery q(m_db);
|
|
int nRemaining = removeCollectionFolder(q, collection, folder_id);
|
|
if (nRemaining == -1) {
|
|
qWarning().nospace() << "Database ERROR: Cannot remove collection " << collection << " from folder "
|
|
<< folder_id << ": " << q.lastError();
|
|
return false;
|
|
}
|
|
removeGuiFolderById(collection, folder_id);
|
|
|
|
if (!sqlPruneCollections(q)) {
|
|
qWarning() << "Database ERROR: Cannot prune collections:" << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
// Keep folder if it is still referenced
|
|
if (nRemaining)
|
|
return true;
|
|
|
|
// Remove the last reference to a folder
|
|
|
|
// First remove all upcoming jobs associated with this folder
|
|
removeFolderFromDocumentQueue(folder_id);
|
|
|
|
// Get a list of all documents associated with folder
|
|
QList<int> documentIds;
|
|
if (!selectDocuments(q, folder_id, &documentIds)) {
|
|
qWarning() << "ERROR: Cannot select documents" << folder_id << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
// Remove all chunks and documents associated with this folder
|
|
for (int document_id: std::as_const(documentIds)) {
|
|
if (!removeChunksByDocumentId(q, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove chunks of document_id" << document_id << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
if (!removeDocument(q, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove document_id" << document_id << q.lastError();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (!removeFolderFromDB(q, folder_id)) {
|
|
qWarning() << "ERROR: Cannot remove folder_id" << folder_id << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
m_collectionMap.remove(folder_id);
|
|
removeFolderFromWatch(path);
|
|
return true;
|
|
}
|
|
|
|
void Database::addFolderToWatch(const QString &path)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "addFolderToWatch" << path;
|
|
#endif
|
|
// pre-check because addPath returns false for already watched paths
|
|
if (!m_watchedPaths.contains(path)) {
|
|
if (!m_watcher->addPath(path))
|
|
qWarning() << "Database::addFolderToWatch: failed to watch" << path;
|
|
// add unconditionally to suppress repeated warnings
|
|
m_watchedPaths << path;
|
|
}
|
|
}
|
|
|
|
void Database::removeFolderFromWatch(const QString &path)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "removeFolderFromWatch" << path;
|
|
#endif
|
|
QDirIterator it(path, QDir::Readable | QDir::Dirs | QDir::NoDotAndDotDot, QDirIterator::Subdirectories);
|
|
QStringList children { path };
|
|
while (it.hasNext())
|
|
children.append(it.next());
|
|
|
|
m_watcher->removePaths(children);
|
|
m_watchedPaths -= QSet(children.begin(), children.end());
|
|
}
|
|
|
|
QList<int> Database::searchEmbeddings(const std::vector<float> &query, const QList<QString> &collections, int nNeighbors)
|
|
{
|
|
constexpr int BATCH_SIZE = 2048;
|
|
|
|
const int n_embd = query.size();
|
|
const us::metric_punned_t metric(n_embd, us::metric_kind_t::ip_k); // inner product
|
|
|
|
QSqlQuery q(m_db);
|
|
if (!q.exec(GET_COLLECTION_EMBEDDINGS_SQL.arg(collections.join("', '")))) {
|
|
qWarning() << "Database ERROR: Failed to exec embeddings query:" << q.lastError();
|
|
return {};
|
|
}
|
|
|
|
us::executor_default_t executor(std::thread::hardware_concurrency());
|
|
us::exact_search_t search;
|
|
|
|
QList<int> batchChunkIds;
|
|
QList<float> batchEmbeddings;
|
|
batchChunkIds.reserve(BATCH_SIZE);
|
|
batchEmbeddings.reserve(BATCH_SIZE * n_embd);
|
|
|
|
struct Result { int chunkId; us::distance_punned_t dist; };
|
|
QList<Result> results;
|
|
|
|
while (q.at() != QSql::AfterLastRow) { // batches
|
|
batchChunkIds.clear();
|
|
batchEmbeddings.clear();
|
|
|
|
while (batchChunkIds.count() < BATCH_SIZE && q.next()) { // batch
|
|
batchChunkIds << q.value(0).toInt();
|
|
batchEmbeddings.resize(batchEmbeddings.size() + n_embd);
|
|
QVariant embdCol = q.value(1);
|
|
if (embdCol.userType() != QMetaType::QByteArray) {
|
|
qWarning() << "Database ERROR: Expected embedding to be blob, got" << embdCol.userType();
|
|
return {};
|
|
}
|
|
auto *embd = static_cast<const QByteArray *>(embdCol.constData());
|
|
const int embd_stride = n_embd * sizeof(float);
|
|
if (embd->size() != embd_stride) {
|
|
qWarning() << "Database ERROR: Expected embedding to be" << embd_stride << "bytes, got"
|
|
<< embd->size();
|
|
return {};
|
|
}
|
|
memcpy(&*(batchEmbeddings.end() - n_embd), embd->constData(), embd_stride);
|
|
}
|
|
|
|
int nBatch = batchChunkIds.count();
|
|
if (!nBatch)
|
|
break;
|
|
|
|
// get top-k nearest neighbors of this batch
|
|
int kBatch = qMin(nNeighbors, nBatch);
|
|
us::exact_search_results_t batchResults = search(
|
|
(us::byte_t const *)batchEmbeddings.data(), nBatch, n_embd * sizeof(float),
|
|
(us::byte_t const *)query.data(), 1, n_embd * sizeof(float),
|
|
kBatch, metric
|
|
);
|
|
|
|
for (int i = 0; i < kBatch; ++i) {
|
|
auto offset = batchResults.at(0)[i].offset;
|
|
us::distance_punned_t distance = batchResults.at(0)[i].distance;
|
|
results.append({batchChunkIds[offset], distance});
|
|
}
|
|
}
|
|
|
|
// get top-k nearest neighbors of combined results
|
|
nNeighbors = qMin(nNeighbors, results.size());
|
|
std::partial_sort(
|
|
results.begin(), results.begin() + nNeighbors, results.end(),
|
|
[](const Result &a, const Result &b) { return a.dist < b.dist; }
|
|
);
|
|
|
|
QList<int> chunkIds;
|
|
chunkIds.reserve(nNeighbors);
|
|
for (int i = 0; i < nNeighbors; i++)
|
|
chunkIds << results[i].chunkId;
|
|
return chunkIds;
|
|
}
|
|
|
|
void Database::retrieveFromDB(const QList<QString> &collections, const QString &text, int retrievalSize,
|
|
QList<ResultInfo> *results)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "retrieveFromDB" << collections << text << retrievalSize;
|
|
#endif
|
|
|
|
std::vector<float> queryEmbd = m_embLLM->generateQueryEmbedding(text);
|
|
if (queryEmbd.empty()) {
|
|
qDebug() << "ERROR: generating embeddings returned a null result";
|
|
return;
|
|
}
|
|
|
|
QList<int> searchResults = searchEmbeddings(queryEmbd, collections, retrievalSize);
|
|
if (searchResults.isEmpty())
|
|
return;
|
|
|
|
QSqlQuery q(m_db);
|
|
if (!selectChunk(q, searchResults, retrievalSize)) {
|
|
qDebug() << "ERROR: selecting chunks:" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
while (q.next()) {
|
|
#if defined(DEBUG)
|
|
const int rowid = q.value(0).toInt();
|
|
#endif
|
|
const QString document_path = q.value(2).toString();
|
|
const QString chunk_text = q.value(3).toString();
|
|
const QString date = QDateTime::fromMSecsSinceEpoch(q.value(1).toLongLong()).toString("yyyy, MMMM dd");
|
|
const QString file = q.value(4).toString();
|
|
const QString title = q.value(5).toString();
|
|
const QString author = q.value(6).toString();
|
|
const int page = q.value(7).toInt();
|
|
const int from = q.value(8).toInt();
|
|
const int to = q.value(9).toInt();
|
|
const QString collectionName = q.value(10).toString();
|
|
ResultInfo info;
|
|
info.collection = collectionName;
|
|
info.path = document_path;
|
|
info.file = file;
|
|
info.title = title;
|
|
info.author = author;
|
|
info.date = date;
|
|
info.text = chunk_text;
|
|
info.page = page;
|
|
info.from = from;
|
|
info.to = to;
|
|
results->append(info);
|
|
#if defined(DEBUG)
|
|
qDebug() << "retrieve rowid:" << rowid
|
|
<< "chunk_text:" << chunk_text;
|
|
#endif
|
|
}
|
|
}
|
|
|
|
// FIXME This is very slow and non-interruptible and when we close the application and we're
|
|
// cleaning a large table this can cause the app to take forever to shut down. This would ideally be
|
|
// interruptible and we'd continue 'cleaning' when we restart
|
|
bool Database::cleanDB()
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "cleanDB";
|
|
#endif
|
|
|
|
// Scan all folders in db to make sure they still exist
|
|
QSqlQuery q(m_db);
|
|
QList<CollectionItem> collections;
|
|
if (!selectAllFromCollections(q, &collections)) {
|
|
qWarning() << "ERROR: Cannot select collections" << q.lastError();
|
|
return false;
|
|
}
|
|
|
|
transaction();
|
|
|
|
for (const auto &i: std::as_const(collections)) {
|
|
// Find the path for the folder
|
|
QFileInfo info(i.folder_path);
|
|
if (!info.exists() || !info.isReadable()) {
|
|
#if defined(DEBUG)
|
|
qDebug() << "clean db removing folder" << i.folder_id << i.folder_path;
|
|
#endif
|
|
if (!removeFolderInternal(i.collection, i.folder_id, i.folder_path)) {
|
|
rollback();
|
|
return false;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Scan all documents in db to make sure they still exist
|
|
if (!q.prepare(SELECT_ALL_DOCUMENTS_SQL)) {
|
|
qWarning() << "ERROR: Cannot prepare sql for select all documents" << q.lastError();
|
|
rollback();
|
|
return false;
|
|
}
|
|
|
|
if (!q.exec()) {
|
|
qWarning() << "ERROR: Cannot exec sql for select all documents" << q.lastError();
|
|
rollback();
|
|
return false;
|
|
}
|
|
|
|
while (q.next()) {
|
|
int document_id = q.value(0).toInt();
|
|
QString document_path = q.value(1).toString();
|
|
QFileInfo info(document_path);
|
|
if (info.exists() && info.isReadable() && m_scannedFileExtensions.contains(info.suffix()))
|
|
continue;
|
|
|
|
#if defined(DEBUG)
|
|
qDebug() << "clean db removing document" << document_id << document_path;
|
|
#endif
|
|
|
|
// Remove all chunks and documents that either don't exist or have become unreadable
|
|
QSqlQuery query(m_db);
|
|
if (!removeChunksByDocumentId(query, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove chunks of document_id" << document_id << query.lastError();
|
|
rollback();
|
|
return false;
|
|
}
|
|
|
|
if (!removeDocument(query, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove document_id" << document_id << query.lastError();
|
|
rollback();
|
|
return false;
|
|
}
|
|
}
|
|
|
|
commit();
|
|
return true;
|
|
}
|
|
|
|
void Database::changeChunkSize(int chunkSize)
|
|
{
|
|
if (chunkSize == m_chunkSize)
|
|
return;
|
|
|
|
#if defined(DEBUG)
|
|
qDebug() << "changeChunkSize" << chunkSize;
|
|
#endif
|
|
|
|
QSqlQuery q(m_db);
|
|
// Scan all documents in db to make sure they still exist
|
|
if (!q.prepare(SELECT_ALL_DOCUMENTS_SQL)) {
|
|
qWarning() << "ERROR: Cannot prepare sql for select all documents" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
if (!q.exec()) {
|
|
qWarning() << "ERROR: Cannot exec sql for select all documents" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
transaction();
|
|
|
|
while (q.next()) {
|
|
int document_id = q.value(0).toInt();
|
|
// Remove all chunks and documents to change the chunk size
|
|
QSqlQuery query(m_db);
|
|
if (!removeChunksByDocumentId(query, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove chunks of document_id" << document_id << query.lastError();
|
|
return rollback();
|
|
}
|
|
|
|
if (!removeDocument(query, document_id)) {
|
|
qWarning() << "ERROR: Cannot remove document_id" << document_id << query.lastError();
|
|
return rollback();
|
|
}
|
|
}
|
|
|
|
commit();
|
|
|
|
m_chunkSize = chunkSize;
|
|
addCurrentFolders();
|
|
updateCollectionStatistics();
|
|
}
|
|
|
|
void Database::changeFileExtensions(const QStringList &extensions)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "changeFileExtensions";
|
|
#endif
|
|
|
|
m_scannedFileExtensions = extensions;
|
|
|
|
if (cleanDB())
|
|
updateCollectionStatistics();
|
|
|
|
QSqlQuery q(m_db);
|
|
QList<CollectionItem> collections;
|
|
if (!selectAllFromCollections(q, &collections)) {
|
|
qWarning() << "ERROR: Cannot select collections" << q.lastError();
|
|
return;
|
|
}
|
|
|
|
for (const auto &i: std::as_const(collections)) {
|
|
if (!i.forceIndexing)
|
|
scanDocuments(i.folder_id, i.folder_path);
|
|
}
|
|
}
|
|
|
|
void Database::directoryChanged(const QString &path)
|
|
{
|
|
#if defined(DEBUG)
|
|
qDebug() << "directoryChanged" << path;
|
|
#endif
|
|
|
|
// search for a collection that contains this folder (we watch subdirectories)
|
|
int folder_id = -1;
|
|
QDir dir(path);
|
|
for (;;) {
|
|
QSqlQuery q(m_db);
|
|
if (!selectFolder(q, dir.path(), &folder_id)) {
|
|
qWarning() << "ERROR: Cannot select folder from path" << dir.path() << q.lastError();
|
|
return;
|
|
}
|
|
if (folder_id != -1)
|
|
break;
|
|
|
|
// check next parent
|
|
if (!dir.cdUp()) {
|
|
if (!dir.isRoot()) break; // folder removed
|
|
Q_ASSERT(false);
|
|
qWarning() << "ERROR: Watched folder does not exist in db" << path;
|
|
m_watcher->removePath(path);
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Clean the database
|
|
if (cleanDB())
|
|
updateCollectionStatistics();
|
|
|
|
// Rescan the documents associated with the folder
|
|
if (folder_id != -1)
|
|
scanDocuments(folder_id, path);
|
|
}
|