Optimizes ChatStreamObserver to use only the last message instead of parsing all messages again

This commit is contained in:
Markus Klenke 2024-02-16 16:31:39 +00:00
parent 0f414f71a3
commit cf4e7a96e8
2 changed files with 3 additions and 6 deletions

View File

@ -2,10 +2,8 @@ package io.github.amithkoujalgi.ollama4j.core.models.chat;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.stream.Collectors;
import io.github.amithkoujalgi.ollama4j.core.OllamaStreamHandler; import io.github.amithkoujalgi.ollama4j.core.OllamaStreamHandler;
import lombok.NonNull;
public class OllamaChatStreamObserver { public class OllamaChatStreamObserver {
@ -13,7 +11,7 @@ public class OllamaChatStreamObserver {
private List<OllamaChatResponseModel> responseParts = new ArrayList<>(); private List<OllamaChatResponseModel> responseParts = new ArrayList<>();
private String message; private String message = "";
public OllamaChatStreamObserver(OllamaStreamHandler streamHandler) { public OllamaChatStreamObserver(OllamaStreamHandler streamHandler) {
this.streamHandler = streamHandler; this.streamHandler = streamHandler;
@ -25,8 +23,7 @@ public class OllamaChatStreamObserver {
} }
protected void handleCurrentResponsePart(OllamaChatResponseModel currentResponsePart){ protected void handleCurrentResponsePart(OllamaChatResponseModel currentResponsePart){
List<@NonNull String> allResponsePartsByNow = responseParts.stream().map(r -> r.getMessage().getContent()).collect(Collectors.toList()); message = message + currentResponsePart.getMessage().getContent();
message = String.join("", allResponsePartsByNow);
streamHandler.accept(message); streamHandler.accept(message);
} }

View File

@ -183,7 +183,7 @@ class TestRealAPIs {
OllamaChatResult chatResult = ollamaAPI.chat(requestModel,(s) -> { OllamaChatResult chatResult = ollamaAPI.chat(requestModel,(s) -> {
LOG.info(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); LOG.info(substring);
sb.append(substring); sb.append(substring);
}); });