Update the sources in chatllm directly and make it thread safe to do so.

Signed-off-by: Adam Treat <treat.adam@gmail.com>
nomic-ai · Nov 4, 2024 · a4f55f6 · a4f55f6
1 parent 50ea10e
commit a4f55f6
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 30 deletions.
diff --git a/gpt4all-chat/qml/ChatItemView.qml b/gpt4all-chat/qml/ChatItemView.qml
@@ -351,6 +351,8 @@ GridLayout {
         Layout.preferredWidth: childrenRect.width
         Layout.preferredHeight: childrenRect.height
         visible: {
+            if (name !== "Response: ")
+                return false
             if (consolidatedSources.length === 0)
                 return false
             if (!MySettings.localDocsShowReferences)

diff --git a/gpt4all-chat/src/chat.cpp b/gpt4all-chat/src/chat.cpp
@@ -368,8 +368,6 @@ QString Chat::fallbackReason() const
 void Chat::handleDatabaseResultsChanged(const QList<ResultInfo> &results)
 {
     m_databaseResults = results;
-    const int index = m_chatModel->count() - 1;
-    m_chatModel->updateSources(index, m_databaseResults);
     m_needsSave = true;
 }
 

diff --git a/gpt4all-chat/src/chatllm.cpp b/gpt4all-chat/src/chatllm.cpp
@@ -776,36 +776,16 @@ auto ChatLLM::promptInternal(const QStringList &enabledCollections, const LLMode
         QString query;
         {
             auto items = m_chatModel->chatItems(); // holds lock
+            Q_ASSERT(items.size() > 1);
             auto prompt = items.end()[-2];
             Q_ASSERT(prompt.name == "Prompt: "_L1);
             query = prompt.value;
         }
         emit requestRetrieveFromDB(enabledCollections, query, retrievalSize, &databaseResults); // blocks
+        m_chatModel->updateSources(databaseResults);
         emit databaseResultsChanged(databaseResults);
     }
 
-    // Augment the prompt template with the results if any
-    QString docsContext;
-    if (!databaseResults.isEmpty()) {
-        QStringList results;
-        results.reserve(databaseResults.size());
-        for (const ResultInfo &info : databaseResults)
-            results << u"Collection: %1\nPath: %2\nExcerpt: %3"_s.arg(info.collection, info.path, info.text);
-
-        // FIXME(jared): use a Jinja prompt template instead of hardcoded Alpaca-style localdocs template
-        docsContext = u"### Context:\n%1\n\n"_s.arg(results.join("\n\n"));
-    }
-
-    // TODO(jared): pass this as a Jinja arg
-#if 0
-    if (!docsContext.isEmpty()) {
-        auto old_n_predict = std::exchange(m_ctx.n_predict, 0); // decode localdocs context without a response
-        m_llModelInfo.model->prompt(docsContext.toStdString(), "%1", promptFunc, responseFunc,
-                                    /*allowContextShift*/ true, m_ctx);
-        m_ctx.n_predict = old_n_predict; // now we are ready for a response
-    }
-#endif
-
     auto [nMessages, conversation] = applyJinjaTemplate();
 
     if (!dynamic_cast<const ChatAPI *>(m_llModelInfo.model.get())) {

diff --git a/gpt4all-chat/src/chatmodel.h b/gpt4all-chat/src/chatmodel.h
@@ -361,16 +361,22 @@ class ChatModel : public QAbstractListModel
         return consolidatedSources;
     }
 
-    Q_INVOKABLE void updateSources(int index, const QList<ResultInfo> &sources)
+    Q_INVOKABLE void updateSources(const QList<ResultInfo> &sources)
     {
         {
             QMutexLocker locker(&m_mutex);
-            if (index < 0 || index >= m_chatItems.size()) return;
-
-            ChatItem &item = m_chatItems[index];
-            item.sources = sources;
-            item.consolidatedSources = consolidateSources(sources);
+            Q_ASSERT(m_chatItems.size() > 1);
+            if (m_chatItems.size() < 2) return;
+            ChatItem &r = m_chatItems.back();
+            Q_ASSERT(r.type() == ChatItem::Type::Response);
+            r.sources = sources;
+            r.consolidatedSources = consolidateSources(sources);
+            ChatItem &p = m_chatItems.end()[-2];
+            Q_ASSERT(p.type() == ChatItem::Type::Prompt);
+            p.sources = sources;
+            p.consolidatedSources = consolidateSources(sources);
         }
+        const int index = m_chatItems.size() - 1;
         emit dataChanged(createIndex(index, 0), createIndex(index, 0), {SourcesRole});
         emit dataChanged(createIndex(index, 0), createIndex(index, 0), {ConsolidatedSourcesRole});
     }