From 1ba555a1742f24c607810748d63c7e12ab4ba10c Mon Sep 17 00:00:00 2001
From: Jared Van Bortel <jared@nomic.ai>
Date: Fri, 28 Feb 2025 12:32:48 -0500
Subject: [PATCH] fix handling of responses that come in chunks, and non-200
 status codes

---
 .../include/gpt4all-backend/ollama-client.h   | 16 ++++----
 gpt4all-backend/src/ollama-client.cpp         | 38 ++++++++++++++++---
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/gpt4all-backend/include/gpt4all-backend/ollama-client.h b/gpt4all-backend/include/gpt4all-backend/ollama-client.h
index da5e448c..4f19e1fe 100644
--- a/gpt4all-backend/include/gpt4all-backend/ollama-client.h
+++ b/gpt4all-backend/include/gpt4all-backend/ollama-client.h
@@ -4,7 +4,7 @@
 
 #include <QCoro/QCoroTask> // IWYU pragma: keep
 #include <boost/json.hpp> // IWYU pragma: keep
-#include <boost/system.hpp>
+#include <boost/system.hpp> // IWYU pragma: keep
 
 #include <QNetworkAccessManager>
 #include <QNetworkReply>
@@ -17,27 +17,27 @@
 #include <variant>
 
 class QNetworkRequest;
+class QRestReply;
 
 
 namespace gpt4all::backend {
 
 struct ResponseError {
+public:
+    struct BadStatus { int code; };
+
 private:
     using ErrorCode = std::variant<
         QNetworkReply::NetworkError,
-        boost::system::error_code
+        boost::system::error_code,
+        BadStatus
     >;
 
 public:
     ErrorCode error;
     QString   errorString;
 
-    ResponseError(const QNetworkReply *reply)
-        : error(reply->error())
-        , errorString(reply->errorString())
-    {
-        assert(reply->error());
-    }
+    ResponseError(const QRestReply *reply);
 
     ResponseError(const boost::system::system_error &e)
         : error(e.code())
diff --git a/gpt4all-backend/src/ollama-client.cpp b/gpt4all-backend/src/ollama-client.cpp
index 5a1f1ddf..f7f0e2e0 100644
--- a/gpt4all-backend/src/ollama-client.cpp
+++ b/gpt4all-backend/src/ollama-client.cpp
@@ -8,6 +8,9 @@
 
 #include <QByteArray>
 #include <QNetworkRequest>
+#include <QRestReply>
+#include <QVariant>
+#include <QtAssert>
 
 #include <coroutine>
 #include <expected>
@@ -21,6 +24,26 @@ namespace json = boost::json;
 namespace gpt4all::backend {
 
 
+ResponseError::ResponseError(const QRestReply *reply)
+{
+    auto *nr = reply->networkReply();
+    if (reply->hasError()) {
+        error       = nr->error();
+        errorString = nr->errorString();
+    } else if (!reply->isHttpStatusSuccess()) {
+        auto code   = reply->httpStatus();
+        auto reason = nr->attribute(QNetworkRequest::HttpReasonPhraseAttribute);
+        error       = BadStatus(code);
+        errorString = u"HTTP %1%2%3 for URL \"%4\""_s.arg(
+            QString::number(code),
+            reason.isValid() ? u" "_s : QString(),
+            reason.toString(),
+            nr->request().url().toString()
+        );
+    } else
+        Q_UNREACHABLE();
+}
+
 QNetworkRequest OllamaClient::makeRequest(const QString &path) const
 {
     QNetworkRequest req(m_baseUrl.resolved(QUrl(path)));
@@ -30,16 +53,21 @@ QNetworkRequest OllamaClient::makeRequest(const QString &path) const
 
 auto OllamaClient::processResponse(QNetworkReply &reply) -> QCoro::Task<DataOrRespErr<json::value>>
 {
+    QRestReply restReply(&reply);
     if (reply.error())
-        co_return std::unexpected(&reply);
+        co_return std::unexpected(&restReply);
 
     auto coroReply = qCoro(reply);
-    do {
+    for (;;) {
         auto chunk = co_await coroReply.readAll();
-        if (reply.error())
-            co_return std::unexpected(&reply);
+        if (!restReply.isSuccess())
+            co_return std::unexpected(&restReply);
+        if (chunk.isEmpty()) {
+            Q_ASSERT(reply.atEnd());
+            break;
+        }
         m_parser.write(chunk.data(), chunk.size());
-    } while (!reply.atEnd());
+    }
 
     m_parser.finish();
     co_return m_parser.release();