Adds the collections to serialize and implement references for localdocs.

gpt4all-chat/chat.cpp

@@ -15,7 +15,6 @@ Chat::Chat(QObject *parent)
     , m_llmodel(new ChatLLM(this))
     , m_isServer(false)
     , m_shouldDeleteLater(false)
-    , m_contextContainsLocalDocs(false)
 {
     connectLLM();
 }
@@ -31,7 +30,6 @@ Chat::Chat(bool isServer, QObject *parent)
     , m_llmodel(new Server(this))
     , m_isServer(true)
     , m_shouldDeleteLater(false)
-    , m_contextContainsLocalDocs(false)
 {
     connectLLM();
 }
@@ -103,7 +101,8 @@ void Chat::prompt(const QString &prompt, const QString &prompt_template, int32_t
     int32_t top_k, float top_p, float temp, int32_t n_batch, float repeat_penalty,
     int32_t repeat_penalty_tokens)
 {
-    m_contextContainsLocalDocs = false;
+    Q_ASSERT(m_results.isEmpty());
+    m_results.clear(); // just in case, but the assert above is important
     m_responseInProgress = true;
     m_responseState = Chat::LocalDocsRetrieval;
     emit responseInProgressChanged();
@@ -116,18 +115,25 @@ void Chat::prompt(const QString &prompt, const QString &prompt_template, int32_t
     m_queuedPrompt.n_batch = n_batch;
     m_queuedPrompt.repeat_penalty = repeat_penalty;
     m_queuedPrompt.repeat_penalty_tokens = repeat_penalty_tokens;
-    LocalDocs::globalInstance()->requestRetrieve(m_collections, prompt);
+    LocalDocs::globalInstance()->requestRetrieve(m_id, m_collections, prompt);
 }
 
-void Chat::handleLocalDocsRetrieved()
+void Chat::handleLocalDocsRetrieved(const QString &uid, const QList<ResultInfo> &results)
 {
+    // If the uid doesn't match, then these are not our results
+    if (uid != m_id)
+        return;
+
+    // Store our results locally
+    m_results = results;
+
+    // Augment the prompt template with the results if any
     QList<QString> augmentedTemplate;
-    QList<QString> results = LocalDocs::globalInstance()->result();
+    if (!m_results.isEmpty())
-    if (!results.isEmpty()) {
         augmentedTemplate.append("### Context:");
-        augmentedTemplate.append(results.join("\n\n"));
+    for (const ResultInfo &info : m_results)
-    }
+        augmentedTemplate.append(info.text);
-    m_contextContainsLocalDocs = !results.isEmpty();
+
     augmentedTemplate.append(m_queuedPrompt.prompt_template);
     emit promptRequested(
         m_queuedPrompt.prompt,
@@ -191,13 +197,48 @@ void Chat::handleModelLoadedChanged()
 
 void Chat::promptProcessing()
 {
-    m_responseState = m_contextContainsLocalDocs ? Chat::LocalDocsProcessing : Chat::PromptProcessing;
+    m_responseState = !m_results.isEmpty() ? Chat::LocalDocsProcessing : Chat::PromptProcessing;
     emit responseStateChanged();
 }
 
 void Chat::responseStopped()
 {
-    m_contextContainsLocalDocs = false;
+    const QString chatResponse = response();
+    QList<QString> finalResponse { chatResponse };
+    int validReferenceNumber = 1;
+    for (const ResultInfo &info : m_results) {
+        if (info.file.isEmpty())
+            continue;
+        if (validReferenceNumber == 1)
+            finalResponse.append(QStringLiteral("---"));
+        QString reference;
+        {
+            QTextStream stream(&reference);
+            stream << (validReferenceNumber++) << ". ";
+            if (!info.title.isEmpty())
+                stream << "\"" << info.title << "\". ";
+            if (!info.author.isEmpty())
+                stream << "By " << info.author << ". ";
+            if (!info.date.isEmpty())
+                stream << "Date: " << info.date << ". ";
+            stream << "In " << info.file << ". ";
+            if (info.page != -1)
+                stream << "Page " << info.page << ". ";
+            if (info.from != -1) {
+                stream << "Lines " << info.from;
+                if (info.to != -1)
+                    stream << "-" << info.to;
+                stream << ". ";
+            }
+        }
+        finalResponse.append(reference);
+    }
+
+    const int index = m_chatModel->count() - 1;
+    m_chatModel->updateValue(index, finalResponse.join("\n"));
+    emit responseChanged();
+
+    m_results.clear();
     m_responseInProgress = false;
     m_responseState = Chat::ResponseStopped;
     emit responseInProgressChanged();
@@ -301,6 +342,8 @@ bool Chat::serialize(QDataStream &stream, int version) const
     stream << m_name;
     stream << m_userName;
     stream << m_savedModelName;
+    if (version > 2)
+        stream << m_collections;
     if (!m_llmodel->serialize(stream, version))
         return false;
     if (!m_chatModel->serialize(stream, version))
@@ -321,6 +364,10 @@ bool Chat::deserialize(QDataStream &stream, int version)
     // unfortunately, we cannot deserialize these
     if (version < 2 && m_savedModelName.contains("gpt4all-j"))
         return false;
+    if (version > 2) {
+        stream >> m_collections;
+        emit collectionListChanged();
+    }
     m_llmodel->setModelName(m_savedModelName);
     if (!m_llmodel->deserialize(stream, version))
         return false;

gpt4all-chat/chat.h

@@ -7,6 +7,7 @@
 
 #include "chatllm.h"
 #include "chatmodel.h"
+#include "database.h"
 #include "server.h"
 
 class Chat : public QObject
@@ -114,7 +115,7 @@ Q_SIGNALS:
     void collectionListChanged();
 
 private Q_SLOTS:
-    void handleLocalDocsRetrieved();
+    void handleLocalDocsRetrieved(const QString &uid, const QList<ResultInfo> &results);
     void handleResponseChanged();
     void handleModelLoadedChanged();
     void promptProcessing();
@@ -141,6 +142,7 @@ private:
     QString m_userName;
     QString m_savedModelName;
     QList<QString> m_collections;
+    QList<ResultInfo> m_results;
     ChatModel *m_chatModel;
     bool m_responseInProgress;
     ResponseState m_responseState;
@@ -148,7 +150,6 @@ private:
     ChatLLM *m_llmodel;
     bool m_isServer;
     bool m_shouldDeleteLater;
-    bool m_contextContainsLocalDocs;
     Prompt m_queuedPrompt;
 };
 

gpt4all-chat/chatlistmodel.cpp

@@ -6,7 +6,7 @@
 #include <QDataStream>
 
 #define CHAT_FORMAT_MAGIC 0xF5D553CC
-#define CHAT_FORMAT_VERSION 2
+#define CHAT_FORMAT_VERSION 3
 
 ChatListModel::ChatListModel(QObject *parent)
     : QAbstractListModel(parent)

gpt4all-chat/database.cpp

@@ -11,12 +11,14 @@
 
 const auto INSERT_CHUNK_SQL = QLatin1String(R"(
     insert into chunks(document_id, chunk_id, chunk_text,
-        embedding_id, embedding_path) values(?, ?, ?, ?, ?);
+        file, title, author, subject, keywords, page, line_from, line_to,
+        embedding_id, embedding_path) values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
     )");
 
 const auto INSERT_CHUNK_FTS_SQL = QLatin1String(R"(
     insert into chunks_fts(document_id, chunk_id, chunk_text,
-        embedding_id, embedding_path) values(?, ?, ?, ?, ?);
+        file, title, author, subject, keywords, page, line_from, line_to,
+        embedding_id, embedding_path) values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
     )");
 
 const auto DELETE_CHUNKS_SQL = QLatin1String(R"(
@@ -29,16 +31,21 @@ const auto DELETE_CHUNKS_FTS_SQL = QLatin1String(R"(
 
 const auto CHUNKS_SQL = QLatin1String(R"(
     create table chunks(document_id integer, chunk_id integer, chunk_text varchar,
+        file varchar, title varchar, author varchar, subject varchar, keywords varchar,
+        page integer, line_from integer, line_to integer,
         embedding_id integer, embedding_path varchar);
     )");
 
 const auto FTS_CHUNKS_SQL = QLatin1String(R"(
     create virtual table chunks_fts using fts5(document_id unindexed, chunk_id unindexed, chunk_text,
+        file, title, author, subject, keywords, page, line_from, line_to,
         embedding_id unindexed, embedding_path unindexed, tokenize="trigram");
     )");
 
 const auto SELECT_SQL = QLatin1String(R"(
-    select chunks_fts.rowid, chunks_fts.document_id, chunks_fts.chunk_text
+    select chunks_fts.rowid, documents.document_time,
+        chunks_fts.chunk_text, chunks_fts.file, chunks_fts.title, chunks_fts.author, chunks_fts.page,
+        chunks_fts.line_from, chunks_fts.line_to
     from chunks_fts
     join documents ON chunks_fts.document_id = documents.id
     join folders ON documents.folder_id = folders.id
@@ -48,8 +55,10 @@ const auto SELECT_SQL = QLatin1String(R"(
     limit %2;
     )");
 
-bool addChunk(QSqlQuery &q, int document_id, int chunk_id, const QString &chunk_text, int embedding_id,
+bool addChunk(QSqlQuery &q, int document_id, int chunk_id, const QString &chunk_text,
-    const QString &embedding_path)
+    const QString &file, const QString &title, const QString &author, const QString &subject, const QString &keywords,
+    int page, int from, int to,
+    int embedding_id, const QString &embedding_path)
 {
     {
         if (!q.prepare(INSERT_CHUNK_SQL))
@@ -57,6 +66,14 @@ bool addChunk(QSqlQuery &q, int document_id, int chunk_id, const QString &chunk_
         q.addBindValue(document_id);
         q.addBindValue(chunk_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(embedding_id);
         q.addBindValue(embedding_path);
         if (!q.exec())
@@ -68,6 +85,14 @@ bool addChunk(QSqlQuery &q, int document_id, int chunk_id, const QString &chunk_
         q.addBindValue(document_id);
         q.addBindValue(chunk_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(embedding_id);
         q.addBindValue(embedding_path);
         if (!q.exec())
@@ -145,19 +170,6 @@ bool selectChunk(QSqlQuery &q, const QList<QString> &collection_names, const QSt
     return true;
 }
 
-void printResults(QSqlQuery &q)
-{
-    while (q.next()) {
-        int rowid = q.value(0).toInt();
-        QString collection_name = q.value(1).toString();
-        QString chunk_text = q.value(2).toString();
-
-        qDebug() << "rowid:" << rowid
-                 << "collection_name:" << collection_name
-                 << "chunk_text:" << chunk_text;
-    }
-}
-
 const auto INSERT_COLLECTION_SQL = QLatin1String(R"(
     insert into collections(collection_name, folder_id) values(?, ?);
     )");
@@ -457,10 +469,18 @@ QSqlError initDb()
     QString chunk_text1 = "This is an example chunk.";
     QString chunk_text2 = "Another example chunk.";
     QString embedding_path = "/example/embeddings/embedding1.bin";
+    QString file = "document1.txt";
+    QString title;
+    QString author;
+    QString subject;
+    QString keywords;
+    int page = -1;
+    int from = -1;
+    int to = -1;;
     int embedding_id = 1;
 
-    if (!addChunk(q, document_id, 1, chunk_text1, embedding_id, embedding_path) ||
+    if (!addChunk(q, document_id, 1, chunk_text1, file, title, author, subject, keywords, page, from, to, embedding_id, embedding_path) ||
-        !addChunk(q, document_id, 2, chunk_text2, embedding_id, embedding_path)) {
+        !addChunk(q, document_id, 2, chunk_text2, file, title, author, subject, keywords, page, from, to, embedding_id, embedding_path)) {
         qDebug() << "Error adding chunks:" << q.lastError().text();
         return q.lastError();
     }
@@ -468,13 +488,10 @@ QSqlError initDb()
     // Perform a search
     QList<QString> collection_names = {collection_name};
     QString search_text = "example";
-    if (!selectChunk(q, collection_names, search_text)) {
+    if (!selectChunk(q, collection_names, search_text, 3)) {
         qDebug() << "Error selecting chunks:" << q.lastError().text();
         return q.lastError();
     }
-
-    // Print the results
-    printResults(q);
 #endif
 
     return QSqlError();
@@ -499,10 +516,13 @@ void Database::handleDocumentErrorAndScheduleNext(const QString &errorMessage,
         QTimer::singleShot(0, this, &Database::scanQueue);
 }
 
-void Database::chunkStream(QTextStream &stream, int document_id)
+void Database::chunkStream(QTextStream &stream, int document_id, const QString &file,
+    const QString &title, const QString &author, const QString &subject, const QString &keywords, int page)
 {
     int chunk_id = 0;
     int charCount = 0;
+    int line_from = -1;
+    int line_to = -1;
     QList<QString> words;
 
     while (!stream.atEnd()) {
@@ -517,6 +537,14 @@ void Database::chunkStream(QTextStream &stream, int document_id)
                 document_id,
                 ++chunk_id,
                 chunk,
+                file,
+                title,
+                author,
+                subject,
+                keywords,
+                page,
+                line_from,
+                line_to,
                 0 /*embedding_id*/,
                 QString() /*embedding_path*/
             )) {
@@ -604,13 +632,18 @@ void Database::scanQueue()
                 document_id, document_path, q.lastError());
             return;
         }
-        QString text;
         for (int i = 0; i < doc.pageCount(); ++i) {
             const QPdfSelection selection = doc.getAllText(i);
-            text.append(selection.text());
+            QString text = selection.text();
-        }
             QTextStream stream(&text);
-        chunkStream(stream, document_id);
+            chunkStream(stream, document_id, 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(),
+                i + 1
+            );
+        }
     } else {
         QFile file(document_path);
         if (!file.open( QIODevice::ReadOnly)) {
@@ -618,7 +651,8 @@ void Database::scanQueue()
                                                       existing_id, document_path, q.lastError());
         }
         QTextStream stream(&file);
-        chunkStream(stream, document_id);
+        chunkStream(stream, document_id, file.fileName(), QString() /*title*/, QString() /*author*/,
+            QString() /*subject*/, QString() /*keywords*/, -1 /*page*/);
         file.close();
     }
     QSqlDatabase::database().commit();
@@ -867,7 +901,7 @@ bool Database::removeFolderFromWatch(const QString &path)
     return true;
 }
 
-void Database::retrieveFromDB(const QList<QString> &collections, const QString &text, int retrievalSize)
+void Database::retrieveFromDB(const QString &uid, const QList<QString> &collections, const QString &text, int retrievalSize)
 {
 #if defined(DEBUG)
     qDebug() << "retrieveFromDB" << collections << text << retrievalSize;
@@ -879,20 +913,33 @@ void Database::retrieveFromDB(const QList<QString> &collections, const QString &
         return;
     }
 
-    QList<QString> results;
+    QList<ResultInfo> results;
     while (q.next()) {
-        int rowid = q.value(0).toInt();
+        const int rowid = q.value(0).toInt();
-        QString collection_name = q.value(1).toString();
+        const QString date = QDateTime::fromMSecsSinceEpoch(q.value(1).toLongLong()).toString("yyyy, MMMM dd");
-        QString chunk_text = q.value(2).toString();
+        const QString chunk_text = q.value(2).toString();
-        results.append(chunk_text);
+        const QString file = q.value(3).toString();
+        const QString title = q.value(4).toString();
+        const QString author = q.value(5).toString();
+        const int page = q.value(6).toInt();
+        const int from =q.value(7).toInt();
+        const int to =q.value(8).toInt();
+        ResultInfo info;
+        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
-                 << "collection_name:" << collection_name
                  << "chunk_text:" << chunk_text;
 #endif
     }
-
+    emit retrieveResult(uid, results);
-    emit retrieveResult(results);
 }
 
 void Database::cleanDB()

gpt4all-chat/database.h

@@ -14,6 +14,17 @@ struct DocumentInfo
     QFileInfo doc;
 };
 
+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
+};
+
 struct CollectionItem {
     QString collection;
     QString folder_path;
@@ -32,13 +43,13 @@ public Q_SLOTS:
     void scanDocuments(int folder_id, const QString &folder_path);
     void addFolder(const QString &collection, const QString &path);
     void removeFolder(const QString &collection, const QString &path);
-    void retrieveFromDB(const QList<QString> &collections, const QString &text, int retrievalSize);
+    void retrieveFromDB(const QString &uid, const QList<QString> &collections, const QString &text, int retrievalSize);
     void cleanDB();
     void changeChunkSize(int chunkSize);
 
 Q_SIGNALS:
     void docsToScanChanged();
-    void retrieveResult(const QList<QString> &result);
+    void retrieveResult(const QString &uid, const QList<ResultInfo> &result);
     void collectionListUpdated(const QList<CollectionItem> &collectionList);
 
 private Q_SLOTS:
@@ -51,14 +62,15 @@ private Q_SLOTS:
 
 private:
     void removeFolderInternal(const QString &collection, int folder_id, const QString &path);
-    void chunkStream(QTextStream &stream, int document_id);
+    void chunkStream(QTextStream &stream, int document_id, const QString &file,
+        const QString &title, const QString &author, const QString &subject, const QString &keywords, int page);
     void handleDocumentErrorAndScheduleNext(const QString &errorMessage,
         int document_id, const QString &document_path, const QSqlError &error);
 
 private:
     int m_chunkSize;
     QQueue<DocumentInfo> m_docsToScan;
-    QList<QString> m_retrieve;
+    QList<ResultInfo> m_retrieve;
     QThread m_dbThread;
     QFileSystemWatcher *m_watcher;
 };

gpt4all-chat/localdocs.cpp

@@ -29,7 +29,7 @@ LocalDocs::LocalDocs()
     connect(this, &LocalDocs::requestChunkSizeChange, m_database,
         &Database::changeChunkSize, Qt::QueuedConnection);
     connect(m_database, &Database::retrieveResult, this,
-        &LocalDocs::handleRetrieveResult, Qt::QueuedConnection);
+        &LocalDocs::receivedResult, Qt::QueuedConnection);
     connect(m_database, &Database::collectionListUpdated,
         m_localDocsModel, &LocalDocsModel::handleCollectionListUpdated, Qt::QueuedConnection);
 }
@@ -49,10 +49,9 @@ void LocalDocs::removeFolder(const QString &collection, const QString &path)
     emit requestRemoveFolder(collection, path);
 }
 
-void LocalDocs::requestRetrieve(const QList<QString> &collections, const QString &text)
+void LocalDocs::requestRetrieve(const QString &uid, const QList<QString> &collections, const QString &text)
 {
-    m_retrieveResult = QList<QString>();
+    emit requestRetrieveFromDB(uid, collections, text, m_retrievalSize);
-    emit requestRetrieveFromDB(collections, text, m_retrievalSize);
 }
 
 int LocalDocs::chunkSize() const
@@ -83,9 +82,3 @@ void LocalDocs::setRetrievalSize(int retrievalSize)
     m_retrievalSize = retrievalSize;
     emit retrievalSizeChanged();
 }
-
-void LocalDocs::handleRetrieveResult(const QList<QString> &result)
-{
-    m_retrieveResult = result;
-    emit receivedResult();
-}

gpt4all-chat/localdocs.h

@@ -20,9 +20,7 @@ public:
 
     Q_INVOKABLE void addFolder(const QString &collection, const QString &path);
     Q_INVOKABLE void removeFolder(const QString &collection, const QString &path);
-
+    void requestRetrieve(const QString &uid, const QList<QString> &collections, const QString &text);
-    QList<QString> result() const { return m_retrieveResult; }
-    void requestRetrieve(const QList<QString> &collections, const QString &text);
 
     int chunkSize() const;
     void setChunkSize(int chunkSize);
@@ -33,22 +31,18 @@ public:
 Q_SIGNALS:
     void requestAddFolder(const QString &collection, const QString &path);
     void requestRemoveFolder(const QString &collection, const QString &path);
-    void requestRetrieveFromDB(const QList<QString> &collections, const QString &text, int N);
+    void requestRetrieveFromDB(const QString &uid, const QList<QString> &collections, const QString &text, int retrievalSize);
     void requestChunkSizeChange(int chunkSize);
-    void receivedResult();
+    void receivedResult(const QString &uid, const QList<ResultInfo> &result);
     void localDocsModelChanged();
     void chunkSizeChanged();
     void retrievalSizeChanged();
 
-private Q_SLOTS:
-    void handleRetrieveResult(const QList<QString> &result);
-
 private:
     int m_chunkSize;
     int m_retrievalSize;
     LocalDocsModel *m_localDocsModel;
     Database *m_database;
-    QList<QString> m_retrieveResult;
 
 private:
     explicit LocalDocs();