From 92618e5084c04e68bb6987b386d7e510151e6a44 Mon Sep 17 00:00:00 2001
From: Amith Koujalgi <amith.koujalgi@razorthink.com>
Date: Tue, 14 May 2024 10:35:55 +0530
Subject: [PATCH] Updated `OllamaChatResponseModel` to include `done_reason`
 field. Refer to the Ollama version:
 https://github.com/ollama/ollama/releases/tag/v0.1.37

---
 .../models/chat/OllamaChatResponseModel.java  |  3 ++-
 .../request/OllamaChatEndpointCaller.java     | 21 +++++++------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatResponseModel.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatResponseModel.java
index 4d0b027..418338f 100644
--- a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatResponseModel.java
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatResponseModel.java
@@ -1,14 +1,15 @@
 package io.github.amithkoujalgi.ollama4j.core.models.chat;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import lombok.Data;
 
 import java.util.List;
-import lombok.Data;
 
 @Data
 public class OllamaChatResponseModel {
     private String model;
     private @JsonProperty("created_at") String createdAt;
+    private @JsonProperty("done_reason") String doneReason;
     private OllamaChatMessage message;
     private boolean done;
     private String error;
diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/request/OllamaChatEndpointCaller.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/request/OllamaChatEndpointCaller.java
index 811ef11..cc6c7f8 100644
--- a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/request/OllamaChatEndpointCaller.java
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/request/OllamaChatEndpointCaller.java
@@ -1,12 +1,6 @@
 package io.github.amithkoujalgi.ollama4j.core.models.request;
 
-import java.io.IOException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import com.fasterxml.jackson.core.JsonProcessingException;
-
 import io.github.amithkoujalgi.ollama4j.core.OllamaStreamHandler;
 import io.github.amithkoujalgi.ollama4j.core.exceptions.OllamaBaseException;
 import io.github.amithkoujalgi.ollama4j.core.models.BasicAuth;
@@ -15,11 +9,15 @@ import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatResponseModel
 import io.github.amithkoujalgi.ollama4j.core.models.chat.OllamaChatStreamObserver;
 import io.github.amithkoujalgi.ollama4j.core.utils.OllamaRequestBody;
 import io.github.amithkoujalgi.ollama4j.core.utils.Utils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
 
 /**
  * Specialization class for requests
  */
-public class OllamaChatEndpointCaller extends OllamaEndpointCaller{
+public class OllamaChatEndpointCaller extends OllamaEndpointCaller {
 
     private static final Logger LOG = LoggerFactory.getLogger(OllamaChatEndpointCaller.class);
 
@@ -39,14 +37,14 @@ public class OllamaChatEndpointCaller extends OllamaEndpointCaller{
         try {
             OllamaChatResponseModel ollamaResponseModel = Utils.getObjectMapper().readValue(line, OllamaChatResponseModel.class);
             responseBuffer.append(ollamaResponseModel.getMessage().getContent());
-            if(streamObserver != null) {
+            if (streamObserver != null) {
                 streamObserver.notify(ollamaResponseModel);
             }
             return ollamaResponseModel.isDone();
         } catch (JsonProcessingException e) {
-            LOG.error("Error parsing the Ollama chat response!",e);
+            LOG.error("Error parsing the Ollama chat response!", e);
             return true;
-        }         
+        }
     }
 
     public OllamaResult call(OllamaRequestBody body, OllamaStreamHandler streamHandler)
@@ -54,7 +52,4 @@ public class OllamaChatEndpointCaller extends OllamaEndpointCaller{
         streamObserver = new OllamaChatStreamObserver(streamHandler);
         return super.callSync(body);
     }
-
-    
- 
 }