2023-05-23 05:13:42 +03:00
|
|
|
#ifndef DATABASE_H
|
|
|
|
#define DATABASE_H
|
|
|
|
|
2024-04-25 20:16:52 +03:00
|
|
|
#include <QElapsedTimer>
|
|
|
|
#include <QFileInfo>
|
|
|
|
#include <QFileSystemWatcher>
|
2023-05-23 05:13:42 +03:00
|
|
|
#include <QObject>
|
|
|
|
#include <QQueue>
|
|
|
|
#include <QThread>
|
2024-04-25 20:16:52 +03:00
|
|
|
#include <QtSql>
|
2023-05-23 05:13:42 +03:00
|
|
|
|
2024-01-22 20:36:01 +03:00
|
|
|
#include "embllm.h"
|
|
|
|
|
2023-10-24 19:13:32 +03:00
|
|
|
class Embeddings;
|
2024-04-25 20:16:52 +03:00
|
|
|
class QTimer;
|
|
|
|
|
2023-05-23 05:13:42 +03:00
|
|
|
struct DocumentInfo
|
|
|
|
{
|
|
|
|
int folder;
|
|
|
|
QFileInfo doc;
|
2023-10-24 19:13:32 +03:00
|
|
|
int currentPage = 0;
|
|
|
|
size_t currentPosition = 0;
|
|
|
|
bool currentlyProcessing = false;
|
|
|
|
bool isPdf() const {
|
|
|
|
return doc.suffix() == QLatin1String("pdf");
|
|
|
|
}
|
2023-05-23 05:13:42 +03:00
|
|
|
};
|
|
|
|
|
2023-05-24 21:49:43 +03:00
|
|
|
struct ResultInfo {
|
|
|
|
QString file; // [Required] The name of the file, but not the full path
|
|
|
|
QString title; // [Optional] The title of the document
|
|
|
|
QString author; // [Optional] The author of the document
|
|
|
|
QString date; // [Required] The creation or the last modification date whichever is latest
|
|
|
|
QString text; // [Required] The text actually used in the augmented context
|
|
|
|
int page = -1; // [Optional] The page where the text was found
|
|
|
|
int from = -1; // [Optional] The line number where the text begins
|
|
|
|
int to = -1; // [Optional] The line number where the text ends
|
|
|
|
};
|
|
|
|
|
2023-05-23 05:13:42 +03:00
|
|
|
struct CollectionItem {
|
|
|
|
QString collection;
|
|
|
|
QString folder_path;
|
|
|
|
int folder_id = -1;
|
2023-06-03 17:08:59 +03:00
|
|
|
bool installed = false;
|
2023-10-24 19:13:32 +03:00
|
|
|
bool indexing = false;
|
2024-01-22 20:36:01 +03:00
|
|
|
QString error;
|
2023-10-24 19:13:32 +03:00
|
|
|
int currentDocsToIndex = 0;
|
|
|
|
int totalDocsToIndex = 0;
|
|
|
|
size_t currentBytesToIndex = 0;
|
|
|
|
size_t totalBytesToIndex = 0;
|
2024-01-22 20:36:01 +03:00
|
|
|
size_t currentEmbeddingsToIndex = 0;
|
|
|
|
size_t totalEmbeddingsToIndex = 0;
|
2023-05-23 05:13:42 +03:00
|
|
|
};
|
|
|
|
Q_DECLARE_METATYPE(CollectionItem)
|
|
|
|
|
|
|
|
class Database : public QObject
|
|
|
|
{
|
|
|
|
Q_OBJECT
|
|
|
|
public:
|
2023-05-24 03:26:31 +03:00
|
|
|
Database(int chunkSize);
|
2023-10-24 19:13:32 +03:00
|
|
|
virtual ~Database();
|
2023-05-23 05:13:42 +03:00
|
|
|
|
|
|
|
public Q_SLOTS:
|
2024-04-25 20:16:52 +03:00
|
|
|
void start();
|
2023-05-23 05:13:42 +03:00
|
|
|
void scanQueue();
|
2024-04-25 20:16:52 +03:00
|
|
|
void scanDocuments(int folder_id, const QString &folder_path, bool isNew);
|
|
|
|
bool addFolder(const QString &collection, const QString &path, bool fromDb);
|
2023-05-23 05:13:42 +03:00
|
|
|
void removeFolder(const QString &collection, const QString &path);
|
2023-06-01 21:13:12 +03:00
|
|
|
void retrieveFromDB(const QList<QString> &collections, const QString &text, int retrievalSize, QList<ResultInfo> *results);
|
2023-05-23 05:13:42 +03:00
|
|
|
void cleanDB();
|
2023-05-24 03:26:31 +03:00
|
|
|
void changeChunkSize(int chunkSize);
|
2023-05-23 05:13:42 +03:00
|
|
|
|
|
|
|
Q_SIGNALS:
|
|
|
|
void docsToScanChanged();
|
2023-10-24 19:13:32 +03:00
|
|
|
void updateInstalled(int folder_id, bool b);
|
|
|
|
void updateIndexing(int folder_id, bool b);
|
2024-01-22 20:36:01 +03:00
|
|
|
void updateError(int folder_id, const QString &error);
|
2023-10-24 19:13:32 +03:00
|
|
|
void updateCurrentDocsToIndex(int folder_id, size_t currentDocsToIndex);
|
|
|
|
void updateTotalDocsToIndex(int folder_id, size_t totalDocsToIndex);
|
|
|
|
void subtractCurrentBytesToIndex(int folder_id, size_t subtractedBytes);
|
|
|
|
void updateCurrentBytesToIndex(int folder_id, size_t currentBytesToIndex);
|
|
|
|
void updateTotalBytesToIndex(int folder_id, size_t totalBytesToIndex);
|
2024-01-22 20:36:01 +03:00
|
|
|
void updateCurrentEmbeddingsToIndex(int folder_id, size_t currentBytesToIndex);
|
|
|
|
void updateTotalEmbeddingsToIndex(int folder_id, size_t totalBytesToIndex);
|
2024-04-25 20:16:52 +03:00
|
|
|
void addCollectionItem(const CollectionItem &item, bool fromDb);
|
2023-10-24 19:13:32 +03:00
|
|
|
void removeFolderById(int folder_id);
|
2023-05-23 05:13:42 +03:00
|
|
|
void collectionListUpdated(const QList<CollectionItem> &collectionList);
|
|
|
|
|
|
|
|
private Q_SLOTS:
|
|
|
|
void directoryChanged(const QString &path);
|
|
|
|
bool addFolderToWatch(const QString &path);
|
|
|
|
bool removeFolderFromWatch(const QString &path);
|
2024-04-25 20:16:52 +03:00
|
|
|
int addCurrentFolders();
|
2024-01-22 20:36:01 +03:00
|
|
|
void handleEmbeddingsGenerated(const QVector<EmbeddingResult> &embeddings);
|
|
|
|
void handleErrorGenerated(int folder_id, const QString &error);
|
2023-05-23 05:13:42 +03:00
|
|
|
|
|
|
|
private:
|
2024-04-25 20:16:52 +03:00
|
|
|
enum class FolderStatus { Started, Embedding, Complete };
|
|
|
|
struct FolderStatusRecord { qint64 startTime; bool isNew; int numDocs, docsChanged, chunksRead; };
|
|
|
|
|
2023-05-23 05:13:42 +03:00
|
|
|
void removeFolderInternal(const QString &collection, int folder_id, const QString &path);
|
2024-01-22 20:36:01 +03:00
|
|
|
size_t chunkStream(QTextStream &stream, int folder_id, int document_id, const QString &file,
|
2023-10-24 19:13:32 +03:00
|
|
|
const QString &title, const QString &author, const QString &subject, const QString &keywords, int page,
|
|
|
|
int maxChunks = -1);
|
|
|
|
void removeEmbeddingsByDocumentId(int document_id);
|
|
|
|
void scheduleNext(int folder_id, size_t countForFolder);
|
|
|
|
void handleDocumentError(const QString &errorMessage,
|
2023-05-23 05:13:42 +03:00
|
|
|
int document_id, const QString &document_path, const QSqlError &error);
|
2023-10-24 19:13:32 +03:00
|
|
|
size_t countOfDocuments(int folder_id) const;
|
|
|
|
size_t countOfBytes(int folder_id) const;
|
|
|
|
DocumentInfo dequeueDocument();
|
|
|
|
void removeFolderFromDocumentQueue(int folder_id);
|
|
|
|
void enqueueDocumentInternal(const DocumentInfo &info, bool prepend = false);
|
|
|
|
void enqueueDocuments(int folder_id, const QVector<DocumentInfo> &infos);
|
2024-04-25 20:16:52 +03:00
|
|
|
void updateIndexingStatus();
|
|
|
|
void updateFolderStatus(int folder_id, FolderStatus status, int numDocs = -1, bool atStart = false, bool isNew = false);
|
2023-05-23 05:13:42 +03:00
|
|
|
|
|
|
|
private:
|
2023-05-24 03:26:31 +03:00
|
|
|
int m_chunkSize;
|
2024-04-25 20:16:52 +03:00
|
|
|
QTimer *m_scanTimer;
|
2023-10-24 19:13:32 +03:00
|
|
|
QMap<int, QQueue<DocumentInfo>> m_docsToScan;
|
2024-04-25 20:16:52 +03:00
|
|
|
QElapsedTimer m_indexingTimer;
|
|
|
|
QMap<int, FolderStatusRecord> m_foldersBeingIndexed;
|
2023-05-24 21:49:43 +03:00
|
|
|
QList<ResultInfo> m_retrieve;
|
2023-05-23 05:13:42 +03:00
|
|
|
QThread m_dbThread;
|
|
|
|
QFileSystemWatcher *m_watcher;
|
2023-10-24 19:13:32 +03:00
|
|
|
EmbeddingLLM *m_embLLM;
|
|
|
|
Embeddings *m_embeddings;
|
2023-05-23 05:13:42 +03:00
|
|
|
};
|
|
|
|
|
|
|
|
#endif // DATABASE_H
|