From c89440cbca1c048892a6249cecce6b21c6f31cde Mon Sep 17 00:00:00 2001
From: Markus Klenke <markusklenke86@gmail.com>
Date: Tue, 13 Feb 2024 17:56:07 +0000
Subject: [PATCH] Adds OllamaStream handling

---
 .../ollama4j/core/OllamaStreamHandler.java    |  7 ++++
 .../models/chat/OllamaChatStreamObserver.java | 34 +++++++++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaStreamHandler.java
 create mode 100644 src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java

diff --git a/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaStreamHandler.java b/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaStreamHandler.java
new file mode 100644
index 0000000..803f393
--- /dev/null
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/OllamaStreamHandler.java
@@ -0,0 +1,7 @@
+package io.github.amithkoujalgi.ollama4j.core;
+
+import java.util.function.Consumer;
+
+public interface OllamaStreamHandler extends Consumer<String>{
+    void accept(String message);
+}
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
new file mode 100644
index 0000000..6a782f4
--- /dev/null
+++ b/src/main/java/io/github/amithkoujalgi/ollama4j/core/models/chat/OllamaChatStreamObserver.java
@@ -0,0 +1,34 @@
+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 {
+
+    private OllamaStreamHandler streamHandler;
+
+    private List<OllamaChatResponseModel> responseParts = new ArrayList<>();
+
+    private String message;
+
+    public OllamaChatStreamObserver(OllamaStreamHandler streamHandler) {
+        this.streamHandler = streamHandler;
+    }
+
+    public void notify(OllamaChatResponseModel currentResponsePart){
+        responseParts.add(currentResponsePart);
+        handleCurrentResponsePart(currentResponsePart);
+    }
+    
+    protected void handleCurrentResponsePart(OllamaChatResponseModel currentResponsePart){
+        List<@NonNull String> allResponsePartsByNow = responseParts.stream().map(r -> r.getMessage().getContent()).collect(Collectors.toList());
+        message = String.join("", allResponsePartsByNow);
+        streamHandler.accept(message);
+    }
+
+
+}