From 6fa5c323b81d1028df447d71de79606b3b2699aa Mon Sep 17 00:00:00 2001
From: amithkoujalgi <koujalgi.amith@gmail.com>
Date: Sun, 31 Aug 2025 15:56:40 +0530
Subject: [PATCH] Refactor stream handler and update test model usage

Simplified ConsoleOutputStreamHandler to print messages directly without substring logic. Updated WithAuth integration test to use the THINKING_MODEL ('gpt-oss:20b') instead of GENERAL_PURPOSE_MODEL ('gemma3:270m') for model pulling and generation.
---
 .../io/github/ollama4j/impl/ConsoleOutputStreamHandler.java | 6 +-----
 .../java/io/github/ollama4j/integrationtests/WithAuth.java  | 5 +++--
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/github/ollama4j/impl/ConsoleOutputStreamHandler.java b/src/main/java/io/github/ollama4j/impl/ConsoleOutputStreamHandler.java
index c9f8e36..d990006 100644
--- a/src/main/java/io/github/ollama4j/impl/ConsoleOutputStreamHandler.java
+++ b/src/main/java/io/github/ollama4j/impl/ConsoleOutputStreamHandler.java
@@ -3,12 +3,8 @@ package io.github.ollama4j.impl;
 import io.github.ollama4j.models.generate.OllamaStreamHandler;
 
 public class ConsoleOutputStreamHandler implements OllamaStreamHandler {
-    private final StringBuffer response = new StringBuffer();
-
     @Override
     public void accept(String message) {
-        String substr = message.substring(response.length());
-        response.append(substr);
-        System.out.print(substr);
+        System.out.print(message);
     }
 }
diff --git a/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java b/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java
index 4b3d0c7..82349dc 100644
--- a/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java
+++ b/src/test/java/io/github/ollama4j/integrationtests/WithAuth.java
@@ -44,6 +44,7 @@ public class WithAuth {
     private static final String BEARER_AUTH_TOKEN = "secret-token";
 
     private static final String GENERAL_PURPOSE_MODEL = "gemma3:270m";
+    private static final String THINKING_MODEL = "gpt-oss:20b";
 
 
     private static OllamaContainer ollama;
@@ -152,7 +153,7 @@ public class WithAuth {
             throws OllamaBaseException, IOException, InterruptedException, URISyntaxException {
         api.setBearerAuth(BEARER_AUTH_TOKEN);
 
-        api.pullModel(GENERAL_PURPOSE_MODEL);
+        api.pullModel(THINKING_MODEL);
 
         String prompt = "The sun is shining brightly and is directly overhead at the zenith, casting my shadow over my foot, so it must be noon.";
 
@@ -169,7 +170,7 @@ public class WithAuth {
         });
         format.put("required", List.of("isNoon"));
 
-        OllamaResult result = api.generate(GENERAL_PURPOSE_MODEL, prompt, format);
+        OllamaResult result = api.generate(THINKING_MODEL, prompt, format);
 
         assertNotNull(result);
         assertNotNull(result.getResponse());