Streaming Client

README.md

@@ -51,7 +51,9 @@ These profiles must be mapped to fixed categories with the following application
 
 The value is the name of the model profile .properties file, for example "llama3.3-70b-azure".
 
-Profiles must be located in the main resources location.
+Agent and completion requests will be invoked with these categories and will load the corresponding model profile.
+
+Profiles must be located in the classpath resources location (e.g. src/main/resources).
 
 Examples can be found in test resources.
 
@@ -75,8 +77,8 @@ The library supports both chat completions and (legacy) text completions. Choose
     modelProvider=OAI_CHAT_COMPLETION
     modelProvider=OAI_TEXT_COMPLETION
 
-Chat completions are recommended for most use cases and easy to use, since the library only needs to send the actual conversation to the server, whereas for text completions, the library needs to create a raw prompt that includes the correct instruction formatting (e.g. special tokens). Text completions are still preferred if you need more control over the raw prompt or if you don't use an instruct model. For text completions, it is required to select the right prompt template. See package .lmlib.prompt for available prompt templates or implement your own PromptStrategy (must be a bean in the spring context).
+Chat completions are recommended for most use cases and easy to use, since the library only needs to send the actual conversation to the server, whereas for text completions, the library needs to create a raw prompt that includes the correct instruction formatting (e.g. special tokens). Text completions are still preferred if you need more control over the raw prompt or if you don't use an instruct model. For text completions, it is required to select the right prompt template. See package .lmlib.prompt for available prompt templates or implement your own PromptStrategy (must be a bean in the spring context). There is also the special NoPromptStrategy ("no") that does not add any special tokens.
-To activate a prompt template, set the following property in your model profile:
+To activate a prompt template, set the following property in your model profile (corresponds to getPromptName()):
 
     prompts=llama3
 
@@ -109,6 +111,28 @@ Example profile: src/test/resources/llama3.3-70b-azure.properties
     	<artifactId>caffeine</artifactId>
     </dependency>
 
+# Usage
+
+1. Setup dependencies, see above
+2. Configure model profiles, see above
+3. Implement a CachedClientRepository if you need your requests cached
+4. For chats or completions, use ClientService service bean
+5. For agentic flows, create an AgentContext and use AgentExecutorService service bean
+
+# Usage Examples
+
+Client service: clientService.getClient(ModelCategory.MEDIUM);
+Client usage details: CompletionClientTest and AzureAiClientTest
+Agentic flow: AgentTest#agenticFlow  
+
+## Completions and Streaming
+
+The library supports both streaming and non-streaming completions.
+
+Streaming methods accept a Consumer to receive the next chunk (direct Client usage) or a callback on the next AgentContext update (Agents). They also return a full Completion object or the final AgentContext at the end of the execution that also contains the statistics.
+
+Non-streaming methods only return the final objects at the end of the execution. This approach is much more resource-friendly and therefore recommended if you don't need to display or process the results in real-time.
+
 # Completions Cache
 
 All language model requests will be cached if you add a CachedClientRepository bean to your context.
@@ -131,20 +155,6 @@ Example:
 
     assertTrue(Grader.assessScoreBinary("Does the following blog post mention the feature " + feature, resultContext.getStringValue("blogPost"), this.clientService));
 
-# Usage
-
-1. Setup dependencies, see above
-2. Configure model profiles, see above
-3. Implement a CachedClientRepository if you need your requests cached
-4. For chats or completions, use ClientService service bean
-5. For agentic flows, create an AgentContext and use AgentExecutorService service bean
-
-# Usage Examples
-
-Text completions: AzureAiClientTest#completion  
-Chat: AzureAiClientTest#chat  
-Agentic flow: AgentTest#agenticFlow  
-
 # Announcement generated by AgentTest#agenticFlow
 
 ## Introducing Advanced Agent Orchestration: Dynamic Conversations & Parallel Processing

src/main/java/de/extio/lmlib/client/Client.java

@@ -1,5 +1,7 @@
 package de.extio.lmlib.client;
 
+import java.util.function.Consumer;
+
 import de.extio.lmlib.profile.ModelCategory;
 import de.extio.lmlib.profile.ModelProfile.ModelProvider;
 
@@ -13,5 +15,7 @@ public interface Client {
 	
 	Completion conversation(ModelCategory modelCategory, Conversation conversation);
 
+	Completion streamConversation(ModelCategory modelCategory, Conversation conversation, Consumer<String> chunkConsumer);
+	
 	ModelProvider getModelProvider();
 }

src/main/java/de/extio/lmlib/client/ClientService.java

@@ -1,6 +1,7 @@
 package de.extio.lmlib.client;
 
 import java.util.List;
+import java.util.function.Consumer;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.cache.annotation.Cacheable;
@@ -67,6 +68,19 @@ public class ClientService {
 			return completion;
 		}
 
+		@Override
+		public Completion streamConversation(final ModelCategory modelCategory, final Conversation conversation_, final Consumer<String> chunkConsumer) {
+			Conversation conversation = conversation_;
+			for (final var interceptor : this.completionInterceptors) {
+				conversation = interceptor.before(modelCategory, conversation);
+			}
+			var completion = this.client.streamConversation(modelCategory, conversation, chunkConsumer);
+			for (final var interceptor : this.completionInterceptors) {
+				completion = interceptor.after(modelCategory, completion);
+			}
+			return completion;
+		}
+		
 		@Override
 		public ModelProvider getModelProvider() {
 			return this.client.getModelProvider();

src/main/java/de/extio/lmlib/client/azureai/AzureAiClient.java

@@ -12,6 +12,7 @@ import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 import org.slf4j.Logger;
@@ -109,6 +110,11 @@ public final class AzureAiClient implements Client {
 		return this.requestCompletions(chat, modelCategory, modelProfile);
 	}
 
+	@Override
+	public Completion streamConversation(final ModelCategory modelCategory, final Conversation conversation, final Consumer<String> chunkConsumer) {
+		throw new UnsupportedOperationException("Not implemented yet");
+	}
+	
 	private List<ChatRequestMessage> createChats(final Conversation conversation, final ModelProfile modelProfile) {
 		final List<ChatRequestMessage> chat = new ArrayList<>();
 		

src/main/java/de/extio/lmlib/client/cached/CachedClient.java

@@ -7,6 +7,7 @@ import java.time.Duration;
 import java.time.OffsetDateTime;
 import java.util.HexFormat;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import de.extio.lmlib.client.Client;
@@ -42,6 +43,15 @@ public class CachedClient implements Client {
 		return this.getCachedResponse(modelCategory, () -> this.client.conversation(modelCategory, conversation), conversation.getConversation().toString());
 	}
 
+	@Override
+	public Completion streamConversation(final ModelCategory modelCategory, final Conversation conversation, final Consumer<String> chunkConsumer) {
+		final var response = this.getCachedResponse(modelCategory, () -> this.client.streamConversation(modelCategory, conversation, chunkConsumer), conversation.getConversation().toString());
+		if (response.statistics().cached()) {
+			chunkConsumer.accept(response.response());
+		}
+		return response;
+	}
+	
 	@Override
 	public ModelProvider getModelProvider() {
 		return this.client.getModelProvider();

src/main/java/de/extio/lmlib/client/oai/completion/AbstractCompletionClient.java

@@ -5,6 +5,7 @@ import java.text.DecimalFormat;
 import java.time.Duration;
 import java.time.Instant;
 import java.time.LocalDateTime;
+import java.util.function.Consumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -14,6 +15,10 @@ import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.web.reactive.function.client.WebClient;
 
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import de.extio.lmlib.client.Client;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
@@ -46,25 +51,51 @@ public abstract class AbstractCompletionClient implements Client, DisposableBean
 	
 	protected final CompletionStatistics totalStatistics = new CompletionStatistics();
 	
+	protected final ObjectMapper objectMapper;
+
+	public AbstractCompletionClient() {
+		this.objectMapper = new ObjectMapper();
+		this.objectMapper.configure(DeserializationFeature.FAIL_ON_IGNORED_PROPERTIES, false);
+		this.objectMapper.configure(DeserializationFeature.FAIL_ON_INVALID_SUBTYPE, false);
+		this.objectMapper.configure(DeserializationFeature.FAIL_ON_MISSING_EXTERNAL_TYPE_ID_PROPERTY, false);
+		this.objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+		this.objectMapper.configure(DeserializationFeature.FAIL_ON_UNRESOLVED_OBJECT_IDS, false);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_BACKSLASH_ESCAPING_ANY_CHARACTER, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_COMMENTS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_MISSING_VALUES, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_NON_NUMERIC_NUMBERS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_TRAILING_COMMA, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_UNQUOTED_CONTROL_CHARS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_YAML_COMMENTS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_LEADING_DECIMAL_POINT_FOR_NUMBERS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_LEADING_PLUS_SIGN_FOR_NUMBERS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_TRAILING_DECIMAL_POINT_FOR_NUMBERS, true);
+		this.objectMapper.configure(JsonParser.Feature.ALLOW_UNQUOTED_FIELD_NAMES, true);
+	}
+
 	@Override
 	public Completion completion(final ModelCategory modelCategory, final String system, final String text) {
 		return this.conversation(modelCategory, Conversation.create(system, text));
 	}
 	
 	@Override
-	public Completion conversation(ModelCategory modelCategory, final Conversation conversation) {
+	public Completion conversation(final ModelCategory modelCategory, final Conversation conversation) {
-		if (modelCategory == null) {
+		return this.streamConversation(modelCategory, conversation, null);
-			modelCategory = ModelCategory.MEDIUM;
+	}
-		}
+	
+	@Override
+	public Completion streamConversation(final ModelCategory modelCategory_, final Conversation conversation, final Consumer<String> chunkConsumer) {
+		final ModelCategory modelCategory = (modelCategory_ == null) ? ModelCategory.MEDIUM : modelCategory_;
 		final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile());
 		if (modelProfile == null || (modelProfile.modelProvider() != ModelProvider.OAI_TEXT_COMPLETION && modelProfile.modelProvider() != ModelProvider.OAI_CHAT_COMPLETION)) {
 			throw new IllegalArgumentException("ModelProfile " + modelCategory.getModelProfile());
 		}
 		
-		return this.requestCompletion(conversation, modelCategory, modelProfile);
+		return this.requestCompletion(conversation, modelCategory, modelProfile, chunkConsumer);
 	}
 	
-	protected abstract Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile);
+	protected abstract Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile, final Consumer<String> chunkConsumer);
 	
 	protected de.extio.lmlib.client.CompletionStatistics createCompletionStatistics(final ModelProfile modelProfile, final LocalDateTime start, final Usage usage, final String prompt, final String response) {
 		int inTokens = 0;

src/main/java/de/extio/lmlib/client/oai/completion/FinishReasons.java

@@ -0,0 +1,11 @@
+package de.extio.lmlib.client.oai.completion;
+
+public final class FinishReasons {
+
+    public static final String FINISH_REASON_STOP = "stop";
+	
+	public static final String FINISH_REASON_LENGTH = "length";
+	
+	public static final String FINISH_REASON_CONTENT_FILTER = "content_filter";
+    
+}

src/main/java/de/extio/lmlib/client/oai/completion/chat/ChatChoice.java

@@ -2,17 +2,11 @@ package de.extio.lmlib.client.oai.completion.chat;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 
-final class ChatChoice {
+class ChatChoice {
 	
-	public static final String FINISH_REASON_STOP = "stop";
+	protected ChatMessage message;
 		
-	public static final String FINISH_REASON_LENGTH = "length";
+	protected String finishReason;
-	
-	public static final String FINISH_REASON_CONTENT_FILTER = "content_filter";
-	
-	private ChatMessage message;
-		
-	private String finishReason;
 	
 	@JsonProperty("message")
 	public ChatMessage getMessage() {
@@ -23,6 +17,15 @@ final class ChatChoice {
 		this.message = message;
 	}
 
+	@JsonProperty("delta")
+	public ChatMessage getDelta() {
+		return message;
+	}
+
+	public void setDelta(final ChatMessage delta) {
+		this.message = delta;
+	}
+
 	@JsonProperty("finish_reason")
 	public String getFinishReason() {
 		return this.finishReason;

src/main/java/de/extio/lmlib/client/oai/completion/chat/ChatCompletionClient.java

@@ -7,20 +7,25 @@ import java.util.ArrayList;
 import java.util.Deque;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpMethod;
 import org.springframework.stereotype.Component;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.CompletionFinishReason;
 import de.extio.lmlib.client.Conversation;
 import de.extio.lmlib.client.Conversation.Turn;
 import de.extio.lmlib.client.oai.completion.AbstractCompletionClient;
+import de.extio.lmlib.client.oai.completion.FinishReasons;
 import de.extio.lmlib.profile.ModelCategory;
 import de.extio.lmlib.profile.ModelProfile;
 import de.extio.lmlib.profile.ModelProfile.ModelProvider;
+import reactor.core.publisher.Mono;
 import reactor.util.retry.Retry;
 
 @Component
@@ -28,13 +33,17 @@ public class ChatCompletionClient extends AbstractCompletionClient {
 	
 	private static final Logger LOGGER = LoggerFactory.getLogger(ChatCompletionClient.class);
 	
+	public ChatCompletionClient() {
+		super();
+	}
+	
 	@Override
 	public ModelProvider getModelProvider() {
 		return ModelProvider.OAI_CHAT_COMPLETION;
 	}
 	
 	@Override
-	protected Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile) {
+	protected Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile, final Consumer<String> chunkConsumer) {
 		final var chat = createChats(conversation, modelProfile);
 		
 		final var request = new ChatCompletionRequest();
@@ -44,18 +53,71 @@ public class ChatCompletionClient extends AbstractCompletionClient {
 		request.setMaxCompletionTokens(modelProfile.maxTokens());
 		request.setTemperature(modelProfile.temperature());
 		request.setTopP(modelProfile.topP());
-		request.setStream(false);
+		if (chunkConsumer == null) {
+			request.setStream(false);
+		}
+		else {
+			request.setStream(true);
+			request.setStreamOptions(new ChatCompletionStreamOptions());
+			request.getStreamOptions().setIncludeUsage(true);
+		}
 		
 		LOGGER.debug("Requesting chat completion at {}", modelProfile.url());
 		final LocalDateTime start = LocalDateTime.now();
 		final var webClient = this.webClientBuilder.baseUrl(modelProfile.url()).build();
+		
 		final var response = webClient
 				.method(HttpMethod.POST)
 				.uri(uriBuilder -> uriBuilder.path("/v1/chat/completions").build())
 				.header("Content-Type", "application/json")
 				.bodyValue(request)
-				.retrieve()
+				.exchangeToMono(clientResponse -> {
-				.bodyToMono(ChatCompletionResponse.class)
+					if (chunkConsumer == null) {
+						return clientResponse.bodyToMono(ChatCompletionResponse.class);
+					}
+					else {
+						final var chatCompletionResponse = new ChatCompletionResponse();
+						chatCompletionResponse.setChoices(List.of(new ChatChoice()));
+						final StringBuilder contenStringBuilder = new StringBuilder();
+						return clientResponse.bodyToFlux(String.class)
+								.doOnNext(buffer -> {
+									try {
+										if (! "[DONE]".equals(buffer)) {
+											final var streamResponse = this.objectMapper.readValue(buffer, ChatCompletionResponse.class);
+											
+											if (streamResponse.getId() != null && !streamResponse.getId().isEmpty()) {
+												chatCompletionResponse.setId(streamResponse.getId());
+											}
+											if (streamResponse.getObject() != null && !streamResponse.getObject().isEmpty()) {
+												chatCompletionResponse.setObject(streamResponse.getObject());
+											}
+											if (streamResponse.getCreated() != null) {
+												chatCompletionResponse.setCreated(streamResponse.getCreated());
+											}
+											if (streamResponse.getUsage() != null) {
+												chatCompletionResponse.setUsage(streamResponse.getUsage());
+											}
+											if (streamResponse.getChoices() != null && !streamResponse.getChoices().isEmpty()) {
+												if (streamResponse.getChoices().getFirst().getFinishReason() != null) {
+													chatCompletionResponse.getChoices().getFirst().setFinishReason(streamResponse.getChoices().getFirst().getFinishReason());
+												}
+												if (streamResponse.getChoices().getFirst().getMessage() != null && streamResponse.getChoices().getFirst().getMessage().getContent() != null && !streamResponse.getChoices().getFirst().getMessage().getContent().isEmpty()) {
+													contenStringBuilder.append(streamResponse.getChoices().getFirst().getMessage().getContent());
+													chunkConsumer.accept(streamResponse.getChoices().getFirst().getMessage().getContent());
+												}
+											}
+										}
+									}
+									catch (final JsonProcessingException e) {
+										LOGGER.warn("Streamed chunk is not parseable", e);
+									}
+								})
+								.then(Mono.<ChatCompletionResponse> fromSupplier(() -> {
+									chatCompletionResponse.getChoices().getFirst().setMessage(new ChatMessage("assistant", contenStringBuilder.toString()));
+									return chatCompletionResponse;
+								}));
+					}
+				})
 				.retryWhen(Retry.backoff(99, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(10)).jitter(0.25d).doAfterRetry(rs -> LOGGER.warn("Request failed " + rs.failure().getMessage())))
 				.block();
 		
@@ -66,8 +128,8 @@ public class ChatCompletionClient extends AbstractCompletionClient {
 		final var choice = response.getChoices().getFirst();
 		final var content = response.getChoices().getFirst().getMessage().getContent();
 		final var finishReason = switch (choice.getFinishReason()) {
-			case ChatChoice.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
+			case FinishReasons.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
-			case ChatChoice.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
+			case FinishReasons.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
 			default -> CompletionFinishReason.DONE;
 		};
 		final var statistics = createCompletionStatistics(modelProfile, start, response.getUsage(), null, content);

src/main/java/de/extio/lmlib/client/oai/completion/chat/ChatCompletionRequest.java

@@ -22,6 +22,8 @@ final class ChatCompletionRequest {
 	
 	private Boolean stream;
 	
+	private ChatCompletionStreamOptions streamOptions;
+	
 	@JsonProperty("model")
 	public String getModel() {
 		return this.model;
@@ -84,4 +86,14 @@ final class ChatCompletionRequest {
 	public void setStream(final Boolean stream) {
 		this.stream = stream;
 	}
+	
+	@JsonProperty("stream_options")
+	public ChatCompletionStreamOptions getStreamOptions() {
+		return streamOptions;
+	}
+	
+	public void setStreamOptions(ChatCompletionStreamOptions streamOptions) {
+		this.streamOptions = streamOptions;
+	}
+	
 }

src/main/java/de/extio/lmlib/client/oai/completion/chat/ChatCompletionStreamOptions.java

@@ -0,0 +1,17 @@
+package de.extio.lmlib.client.oai.completion.chat;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class ChatCompletionStreamOptions {
+
+    private boolean includeUsage;
+
+    @JsonProperty("include_usage")
+    public boolean isIncludeUsage() {
+        return includeUsage;
+    }
+
+    public void setIncludeUsage(final boolean includeUsage) {
+        this.includeUsage = includeUsage;
+    }
+}

src/main/java/de/extio/lmlib/client/oai/completion/text/Choice.java

@@ -4,12 +4,6 @@ import com.fasterxml.jackson.annotation.JsonProperty;
 
 final class Choice {
 	
-	public static final String FINISH_REASON_STOP = "stop";
-	
-	public static final String FINISH_REASON_LENGTH = "length";
-	
-	public static final String FINISH_REASON_CONTENT_FILTER = "content_filter";
-	
 	private String text;
 	
 	private Double score;

src/main/java/de/extio/lmlib/client/oai/completion/text/CompletionRequest.java

@@ -16,6 +16,8 @@ final class CompletionRequest {
 	
 	private Boolean stream;
 	
+	private CompletionStreamOptions streamOptions;
+	
 	@JsonProperty("model")
 	public String getModel() {
 		return this.model;
@@ -69,4 +71,13 @@ final class CompletionRequest {
 	public void setStream(final Boolean stream) {
 		this.stream = stream;
 	}
+	
+	@JsonProperty("stream_options")
+	public CompletionStreamOptions getStreamOptions() {
+		return streamOptions;
+	}
+	
+	public void setStreamOptions(final CompletionStreamOptions streamOptions) {
+		this.streamOptions = streamOptions;
+	}
 }

src/main/java/de/extio/lmlib/client/oai/completion/text/CompletionResponse.java

@@ -18,7 +18,7 @@ final class CompletionResponse {
 	
 	private String content;
 	
-	private boolean isStoppedEos;
+	private Boolean isStoppedEos;
 
 	private Usage usage;
 
@@ -68,11 +68,11 @@ final class CompletionResponse {
 	}
 	
 	@JsonProperty("stopped_eos")
-	public boolean isStoppedEos() {
+	public Boolean isStoppedEos() {
 		return this.isStoppedEos;
 	}
 	
-	public void setStoppedEos(final boolean isStoppedEos) {
+	public void setStoppedEos(final Boolean isStoppedEos) {
 		this.isStoppedEos = isStoppedEos;
 	}
 		

src/main/java/de/extio/lmlib/client/oai/completion/text/CompletionStreamOptions.java

@@ -0,0 +1,17 @@
+package de.extio.lmlib.client.oai.completion.text;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class CompletionStreamOptions {
+
+    private boolean includeUsage;
+
+    @JsonProperty("include_usage")
+    public boolean isIncludeUsage() {
+        return includeUsage;
+    }
+
+    public void setIncludeUsage(final boolean includeUsage) {
+        this.includeUsage = includeUsage;
+    }
+}

src/main/java/de/extio/lmlib/client/oai/completion/text/TextCompletionClient.java

@@ -5,6 +5,7 @@ import java.time.LocalDateTime;
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.List;
+import java.util.function.Consumer;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -12,16 +13,21 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpMethod;
 import org.springframework.stereotype.Component;
 
+import com.azure.ai.inference.models.ChatChoice;
+import com.fasterxml.jackson.core.JsonProcessingException;
+
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.CompletionFinishReason;
 import de.extio.lmlib.client.Conversation;
 import de.extio.lmlib.client.Conversation.Turn;
 import de.extio.lmlib.client.oai.completion.AbstractCompletionClient;
+import de.extio.lmlib.client.oai.completion.FinishReasons;
 import de.extio.lmlib.profile.ModelCategory;
 import de.extio.lmlib.profile.ModelProfile;
 import de.extio.lmlib.profile.ModelProfile.ModelProvider;
 import de.extio.lmlib.prompt.PromptStrategy;
 import de.extio.lmlib.prompt.PromptStrategyFactory;
+import reactor.core.publisher.Mono;
 import reactor.util.retry.Retry;
 
 @Component
@@ -38,7 +44,7 @@ public class TextCompletionClient extends AbstractCompletionClient {
 	}
 	
 	@Override
-	protected Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile) {
+	protected Completion requestCompletion(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile, final Consumer<String> chunkConsumer) {
 		final var promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
 		if (promptStrategy == null) {
 			throw new IllegalArgumentException("Prompt strategy not found: " + modelProfile.prompt());
@@ -52,7 +58,14 @@ public class TextCompletionClient extends AbstractCompletionClient {
 		request.setMaxTokens(modelProfile.maxTokens());
 		request.setTemperature(modelProfile.temperature());
 		request.setTopP(modelProfile.topP());
-		request.setStream(false);
+		if (chunkConsumer == null) {
+			request.setStream(false);
+		}
+		else {
+			request.setStream(true);
+			request.setStreamOptions(new CompletionStreamOptions());
+			request.getStreamOptions().setIncludeUsage(true);
+		}
 		
 		LOGGER.debug("Requesting completion at {}", modelProfile.url());
 		final LocalDateTime start = LocalDateTime.now();
@@ -62,29 +75,94 @@ public class TextCompletionClient extends AbstractCompletionClient {
 				.uri(uriBuilder -> uriBuilder.path("/v1/completions").build())
 				.header("Content-Type", "application/json")
 				.bodyValue(request)
-				.retrieve()
+				.exchangeToMono(clientResponse -> {
-				.bodyToMono(CompletionResponse.class)
+					if (chunkConsumer == null) {
+						return clientResponse.bodyToMono(CompletionResponse.class);
+					}
+					else {
+						final var CompletionResponse = new CompletionResponse();
+						CompletionResponse.setChoices(List.of(new Choice()));
+						final StringBuilder contentStringBuilder = new StringBuilder();
+						return clientResponse.bodyToFlux(String.class)
+								.doOnNext(buffer -> {
+									try {
+										if (!"[DONE]".equals(buffer)) {
+											final var streamResponse = this.objectMapper.readValue(buffer, CompletionResponse.class);
+											
+											if (streamResponse.getId() != null && !streamResponse.getId().isEmpty()) {
+												CompletionResponse.setId(streamResponse.getId());
+											}
+											if (streamResponse.getObject() != null && !streamResponse.getObject().isEmpty()) {
+												CompletionResponse.setObject(streamResponse.getObject());
+											}
+											if (streamResponse.getCreated() != null) {
+												CompletionResponse.setCreated(streamResponse.getCreated());
+											}
+											if (streamResponse.isStoppedEos() != null) {
+												CompletionResponse.setStoppedEos(streamResponse.isStoppedEos());
+											}
+											if (streamResponse.getContent() != null && !streamResponse.getContent().isEmpty()) {
+												contentStringBuilder.append(streamResponse.getContent());
+											}
+											if (streamResponse.getUsage() != null) {
+												CompletionResponse.setUsage(streamResponse.getUsage());
+											}
+											if (streamResponse.getChoices() != null && !streamResponse.getChoices().isEmpty()) {
+												if (streamResponse.getChoices().getFirst().getFinishReason() != null) {
+													CompletionResponse.getChoices().getFirst().setFinishReason(streamResponse.getChoices().getFirst().getFinishReason());
+												}
+												if (streamResponse.getChoices().getFirst().getScore() != null) {
+													CompletionResponse.getChoices().getFirst().setScore(streamResponse.getChoices().getFirst().getScore());
+												}
+												if (streamResponse.getChoices().getFirst().getText() != null && !streamResponse.getChoices().getFirst().getText().isEmpty()) {
+													contentStringBuilder.append(streamResponse.getChoices().getFirst().getText());
+													chunkConsumer.accept(streamResponse.getChoices().getFirst().getText());
+												}
+											}
+										}
+									}
+									catch (final JsonProcessingException e) {
+										LOGGER.warn("Streamed chunk is not parseable", e);
+									}
+								})
+								.then(Mono.<CompletionResponse> fromSupplier(() -> {
+									CompletionResponse.getChoices().getFirst().setText(contentStringBuilder.toString());
+									return CompletionResponse;
+								}));
+					}
+				})
 				.retryWhen(Retry.backoff(99, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(10)).jitter(0.25d).doAfterRetry(rs -> LOGGER.warn("Request failed " + rs.failure().getMessage())))
 				.block();
 		
-		String content;
+		String content = null;
-		CompletionFinishReason finishReason;
+		CompletionFinishReason finishReason = null;
 		if (response.getChoices() != null && !response.getChoices().isEmpty()) {
 			final var choice = response.getChoices().getFirst();
-			content = promptStrategy.removeEOT(choice.getText());
+			if (response.getChoices().getFirst().getText() != null) {
-			finishReason = switch (choice.getFinishReason()) {
+				content = promptStrategy.removeEOT(choice.getText());
-				case Choice.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
+			}
-				case Choice.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
+			if (response.getChoices().getFirst().getFinishReason() != null) {
-				default -> CompletionFinishReason.DONE;
+				finishReason = switch (choice.getFinishReason()) {
-			};
+					case FinishReasons.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
+					case FinishReasons.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
+					default -> CompletionFinishReason.DONE;
+				};
+			}
 		}
-		else if (response.getContent() != null) {
+		if (content == null && response.getContent() != null) {
 			content = promptStrategy.removeEOT(response.getContent());
-			finishReason = response.isStoppedEos() ? CompletionFinishReason.DONE : CompletionFinishReason.TOKEN_LIMIT_REACHED;
 		}
-		else {
+		if (content == null) {
 			throw new IllegalStateException("No parseable response content from server");
 		}
+		if (finishReason == null) {
+			if (response.isStoppedEos() != null) {
+				finishReason = Boolean.TRUE.equals(response.isStoppedEos()) ? CompletionFinishReason.DONE : CompletionFinishReason.TOKEN_LIMIT_REACHED;
+			}
+			else {
+				finishReason = CompletionFinishReason.DONE;
+			}
+		}
 		
 		final var statistics = createCompletionStatistics(modelProfile, start, response.getUsage(), prompt, response.getContent());
 		return new Completion(content, finishReason, statistics);

src/main/java/de/extio/lmlib/prompt/GemmaPromptStrategy.java

@@ -7,7 +7,7 @@ import org.springframework.stereotype.Component;
  * @see https://ai.google.dev/gemma/docs/core/prompt-structure
  */
 @Component
-public class Gemma2PromptStrategy implements PromptStrategy {
+public class GemmaPromptStrategy implements PromptStrategy {
 	
 	@Override
 	public StringBuilder start(final String system, final String user) {
@@ -37,7 +37,7 @@ public class Gemma2PromptStrategy implements PromptStrategy {
 	
 	@Override
 	public String getPromptName() {
-		return "gemma2";
+		return "gemma";
 	}
 	
 }

src/main/java/de/extio/lmlib/prompt/NoPromptStrategy.java

@@ -0,0 +1,40 @@
+package de.extio.lmlib.prompt;
+
+import org.springframework.stereotype.Component;
+
+@Component
+public class NoPromptStrategy implements PromptStrategy {
+    
+    @Override
+    public String getPromptName() {
+        return "noop";
+    }
+    
+    @Override
+    public StringBuilder start(final String system, final String user) {
+        final var sb = new StringBuilder();
+        if (system != null && !system.isBlank()) {
+            sb.append(system);
+        }
+        if (system != null && !system.isBlank() && user != null && !user.isBlank()) {
+            sb.append("\n");
+        }
+        if (user != null && !user.isBlank()) {
+            sb.append(user);
+        }
+        return sb;
+    }
+    
+    @Override
+    public void next(final StringBuilder prompt, final String assistant, final String user) {
+        prompt.append(assistant);
+        prompt.append("\n");
+        prompt.append(user);
+    }
+    
+    @Override
+    public String removeEOT(final String prompt) {
+        return prompt;
+    }
+    
+}

src/test/java/de/extio/lmlib/client/cached/CachedClientTest.java

@@ -23,7 +23,7 @@ import de.extio.lmlib.client.Conversation.Turn;
 import de.extio.lmlib.client.Conversation.TurnType;
 import de.extio.lmlib.profile.ModelCategory;
 
-//@Disabled("This test requires a running Llama server or an Azure subscription (setup azure key in model profile)")
+@Disabled("This test requires a running Llama server or an Azure subscription (setup azure key in model profile)")
 @SpringBootTest(webEnvironment = WebEnvironment.NONE)
 @SpringBootConfiguration
 @ComponentScan(basePackages = "de.extio.lmlib")
@@ -45,8 +45,8 @@ public class CachedClientTest {
 	
 	@Test
 	void cacheCompletion() throws Exception {
-		var completion = this.clientService.getClient(ModelCategory.MEDIUM).completion(
+		var completion = this.clientService.getClient(ModelCategory.SMALL).completion(
-				ModelCategory.MEDIUM,
+				ModelCategory.SMALL,
 				"You are a funny person",
 				"Tell me a cat joke");
 		LOGGER.info(completion.toString());
@@ -55,8 +55,8 @@ public class CachedClientTest {
 		final var firstResponse = completion.response();
 		assertFalse(firstResponse.isBlank());
 		
-		completion = this.clientService.getClient(ModelCategory.MEDIUM).completion(
+		completion = this.clientService.getClient(ModelCategory.SMALL).completion(
-				ModelCategory.MEDIUM,
+				ModelCategory.SMALL,
 				"You are a funny person",
 				"Tell me a cat joke");
 		LOGGER.info(completion.toString());
@@ -64,8 +64,8 @@ public class CachedClientTest {
 		assertEquals(0, completion.statistics().requests());
 		assertEquals(firstResponse, completion.response());
 		
-		completion = this.clientService.getClient(ModelCategory.MEDIUM).completion(
+		completion = this.clientService.getClient(ModelCategory.SMALL).completion(
-				ModelCategory.MEDIUM,
+				ModelCategory.SMALL,
 				"You are a funny person",
 				"Tell me a dog joke");
 		LOGGER.info(completion.toString());
@@ -100,4 +100,31 @@ public class CachedClientTest {
 		assertEquals(1, completion.statistics().requests());
 		assertNotEquals(firstResponse, completion.response());
 	}
+
+	@Test
+	void cacheStream() throws Exception {
+		final var conversation = Conversation.create("You are an assistant who talks extremely like a pirate", "Hi there, I would like to know the climate in the red sea.");
+		
+		final var sb = new java.lang.StringBuilder();
+		var completion = this.clientService.getClient(ModelCategory.MEDIUM).streamConversation(ModelCategory.MEDIUM, conversation, chunk -> {
+			sb.append(chunk);
+			System.out.print(chunk);
+		});
+		assertFalse(completion.statistics().cached());
+		assertEquals(1, completion.statistics().requests());
+		assertTrue(!sb.isEmpty());
+		assertEquals(sb.toString(), completion.response());
+		
+		System.out.println();
+		LOGGER.info("### Next request");
+
+		final var sb2 = new java.lang.StringBuilder();
+		completion = this.clientService.getClient(ModelCategory.MEDIUM).streamConversation(ModelCategory.MEDIUM, conversation, chunk -> {
+			sb2.append(chunk);
+			System.out.print(chunk);
+		});
+		assertTrue(completion.statistics().cached());
+		assertEquals(sb.toString(), sb2.toString());
+		assertEquals(sb2.toString(), completion.response());
+	}
 }

src/test/java/de/extio/lmlib/client/oai/CompletionClientTest.java

@@ -1,5 +1,6 @@
 package de.extio.lmlib.client.oai;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.time.LocalDateTime;
@@ -50,7 +51,7 @@ public class CompletionClientTest {
 		
 		LOGGER.info(completion.response());
 		
-		assertTrue(Grader.assessScoreBinary("Does the text mention the color green?", completion.response(), this.textCompletionClient));
+		assertTrue(Grader.assessScoreBinary("Does the text mention the color green?", completion.response(), this.chatCompletionclient));
 	}
 	
 	@Test
@@ -58,11 +59,26 @@ public class CompletionClientTest {
 		final var conversation = Conversation.create("You are a helpful assistant", "I would like to paint my canvas in a solid color.");
 		conversation.addTurn(new Turn(TurnType.ASSISTANT, "Which color do you want to see?"));
 		conversation.addTurn(new Turn(TurnType.USER, "You choose!"));
-		final var completion = this.textCompletionClient.conversation(ModelCategory.MEDIUM, conversation);
+		final var completion = this.textCompletionClient.conversation(ModelCategory.SMALL, conversation);
 		
 		LOGGER.info(completion.response());
 		
-		assertTrue(Grader.assessScoreBinary("Does the text mention a color?", completion.response(), this.textCompletionClient));
+		assertTrue(Grader.assessScoreBinary("Does the text mention a color?", completion.response(), this.chatCompletionclient));
+	}
+	
+	@Test
+	void streamTextCompletion() throws Exception {
+		final var conversation = Conversation.create("You are a helpful assistant", "I would like to play a role-playing game with you.");
+		conversation.addTurn(new Turn(TurnType.ASSISTANT, "Which character do you choose?"));
+		conversation.addTurn(new Turn(TurnType.USER, "You choose! But I need you to explain the character attributes to me."));
+		final StringBuilder sb = new StringBuilder();
+		final var completion = this.textCompletionClient.streamConversation(ModelCategory.SMALL, conversation, chunk -> {
+			System.out.print(chunk);
+			sb.append(chunk);
+		});
+		
+		assertEquals(completion.response(), sb.toString());
+		assertTrue(Grader.assessScoreBinary("Does the text mention a role-playing game character?", completion.response(), this.chatCompletionclient));
 	}
 	
 	@Test
@@ -74,7 +90,22 @@ public class CompletionClientTest {
 		
 		LOGGER.info(completion.response());
 		
-		assertTrue(Grader.assessScoreBinary("Does the text mention a color?", completion.response(), this.textCompletionClient));
+		assertTrue(Grader.assessScoreBinary("Does the text mention a color?", completion.response(), this.chatCompletionclient));
+	}
+	
+	@Test
+	void streamChatCompletion() throws Exception {
+		final var conversation = Conversation.create("You are a helpful assistant", "I would like to play a game of chess.");
+		conversation.addTurn(new Turn(TurnType.ASSISTANT, "Which opening do you choose?"));
+		conversation.addTurn(new Turn(TurnType.USER, "You choose! But I need you to explain the opening to me."));
+		final StringBuilder sb = new StringBuilder();
+		final var completion = this.chatCompletionclient.streamConversation(ModelCategory.MEDIUM, conversation, chunk -> {
+			System.out.print(chunk);
+			sb.append(chunk);
+		});
+		
+		assertEquals(completion.response(), sb.toString());
+		assertTrue(Grader.assessScoreBinary("Does the text explain a chess opening?", completion.response(), this.chatCompletionclient));
 	}
 	
 	@Disabled

src/test/resources/application-test.properties

@@ -1,8 +1,8 @@
 profile.model.s=local-completion
-profile.model.m=local-completion
+profile.model.m=local-chatcompletion
-profile.model.l=local-completion
+profile.model.l=local-chatcompletion
-profile.model.hot=local-completion
+profile.model.hot=local-chatcompletion
-profile.model.cold=local-completion
+profile.model.cold=local-chatcompletion
 
 agent.threads=2
 

src/test/resources/local-chatcompletion.properties

@@ -0,0 +1,12 @@
+prompts=
+tokenEncoding=cl100k_base
+maxTokens=2000
+maxContextLength=8000
+temperature=0.4
+topP=1.0
+modelProvider=OAI_CHAT_COMPLETION
+modelName=
+url=http://localhost:5001
+apiKey=
+cost1MInTokens=0
+cost1MOutTokens=0

src/test/resources/local-completion.properties

@@ -1,10 +1,10 @@
-prompts=
+prompts=gemma
 tokenEncoding=cl100k_base
 maxTokens=2000
 maxContextLength=8000
 temperature=0.4
 topP=1.0
-modelProvider=OAI_CHAT_COMPLETION
+modelProvider=OAI_TEXT_COMPLETION
 modelName=
 url=http://localhost:5001
 apiKey=