From cf4e7a96e8090d3288fbb686aaf259156f6f7234 Mon Sep 17 00:00:00 2001
From: Markus Klenke <markusklenke86@gmail.com>
Date: Fri, 16 Feb 2024 16:31:39 +0000
Subject: [PATCH] Optimizes ChatStreamObserver to use only the last message
 instead of parsing all messages again

---
 .../core/models/chat/OllamaChatStreamObserver.java         | 7 ++-----
 .../ollama4j/integrationtests/TestRealAPIs.java            | 2 +-
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java
index 6a782f4..ea4b4d8 100644
--- a/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java
@@ -2,10 +2,8 @@ package io.github.amithkoujalgi.ollama4j.core.models.chat;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.Collectors;
 
 import io.github.amithkoujalgi.ollama4j.core.OllamaStreamHandler;
-import lombok.NonNull;
 
 public class OllamaChatStreamObserver {
 
@@ -13,7 +11,7 @@ public class OllamaChatStreamObserver {
 
     private List<OllamaChatResponseModel> responseParts = new ArrayList<>();
 
-    private String message;
+    private String message = "";
 
     public OllamaChatStreamObserver(OllamaStreamHandler streamHandler) {
         this.streamHandler = streamHandler;
@@ -25,8 +23,7 @@ public class OllamaChatStreamObserver {
     }
     
     protected void handleCurrentResponsePart(OllamaChatResponseModel currentResponsePart){
-        List<@NonNull String> allResponsePartsByNow = responseParts.stream().map(r -> r.getMessage().getContent()).collect(Collectors.toList());
-        message = String.join("", allResponsePartsByNow);
+        message = message + currentResponsePart.getMessage().getContent();
         streamHandler.accept(message);
     }
 
diff --git a/src/test/java/io/github/amithkoujalgi/ollama4j/integrationtests/TestRealAPIs.java b/src/test/java/io/github/amithkoujalgi/ollama4j/integrationtests/TestRealAPIs.java
index 870e17f..b4a7e1d 100644
--- a/src/test/java/io/github/amithkoujalgi/ollama4j/integrationtests/TestRealAPIs.java
+++ b/src/test/java/io/github/amithkoujalgi/ollama4j/integrationtests/TestRealAPIs.java
@@ -183,7 +183,7 @@ class TestRealAPIs {
 
       OllamaChatResult chatResult = ollamaAPI.chat(requestModel,(s) -> {
         LOG.info(s);
-        String substring = s.substring(sb.toString().length(), s.length()-1);
+        String substring = s.substring(sb.toString().length(), s.length());
         LOG.info(substring);
         sb.append(substring);
       });