exo/lm-lib
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Streaming Client

Browse Source
This commit is contained in:
extio-de
2025-06-28 14:36:26 +02:00
parent 2350740136
commit 365043b5a0
23 changed files with 473 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
README.md
View File

@@ -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).
To activate a prompt template, set the following property in your model profile:
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 (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

4
src/main/java/de/extio/lmlib/client/Client.java
View File

@@ -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;
@@ -12,6 +14,8 @@ public interface Client {
Completion completion(ModelCategory modelCategory, String system, String text);
Completion conversation(ModelCategory modelCategory, Conversation conversation);
Completion streamConversation(ModelCategory modelCategory, Conversation conversation, Consumer<String> chunkConsumer);
ModelProvider getModelProvider();
}

14
src/main/java/de/extio/lmlib/client/ClientService.java
View File

@@ -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;
@@ -66,6 +67,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() {

6
src/main/java/de/extio/lmlib/client/azureai/AzureAiClient.java
View File

@@ -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;
@@ -108,6 +109,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<>();

10
src/main/java/de/extio/lmlib/client/cached/CachedClient.java
View File

@@ -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;
@@ -41,6 +42,15 @@ public class CachedClient implements Client {
public Completion conversation(final ModelCategory modelCategory, final Conversation conversation) {
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() {

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

@@ -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) {
if (modelCategory == null) {
modelCategory = ModelCategory.MEDIUM;
}
public Completion conversation(final ModelCategory modelCategory, final Conversation conversation) {
return this.streamConversation(modelCategory, conversation, null);
}
@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;

11
src/main/java/de/extio/lmlib/client/oai/completion/FinishReasons.java Normal file
View File

@@ -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";
}

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

@@ -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";
public static final String FINISH_REASON_LENGTH = "length";
public static final String FINISH_REASON_CONTENT_FILTER = "content_filter";
private ChatMessage message;
protected ChatMessage message;
private String finishReason;
protected 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;

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

@@ -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()
.bodyToMono(ChatCompletionResponse.class)
.exchangeToMono(clientResponse -> {
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 ChatChoice.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
case FinishReasons.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
case FinishReasons.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
default -> CompletionFinishReason.DONE;
};
final var statistics = createCompletionStatistics(modelProfile, start, response.getUsage(), null, content);

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

@@ -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;
}
}

17
src/main/java/de/extio/lmlib/client/oai/completion/chat/ChatCompletionStreamOptions.java Normal file
View File

@@ -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;
}
}

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

@@ -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;

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

@@ -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;
}
}

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

@@ -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;
}

17
src/main/java/de/extio/lmlib/client/oai/completion/text/CompletionStreamOptions.java Normal file
View File

@@ -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;
}
}

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

@@ -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()
.bodyToMono(CompletionResponse.class)
.exchangeToMono(clientResponse -> {
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;
CompletionFinishReason finishReason;
String content = null;
CompletionFinishReason finishReason = null;
if (response.getChoices() != null && !response.getChoices().isEmpty()) {
final var choice = response.getChoices().getFirst();
content = promptStrategy.removeEOT(choice.getText());
finishReason = switch (choice.getFinishReason()) {
case Choice.FINISH_REASON_CONTENT_FILTER -> CompletionFinishReason.CONTENT_FILTERED;
case Choice.FINISH_REASON_LENGTH -> CompletionFinishReason.TOKEN_LIMIT_REACHED;
default -> CompletionFinishReason.DONE;
};
if (response.getChoices().getFirst().getText() != null) {
content = promptStrategy.removeEOT(choice.getText());
}
if (response.getChoices().getFirst().getFinishReason() != null) {
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);

4
src/main/java/de/extio/lmlib/prompt/Gemma2PromptStrategy.java → src/main/java/de/extio/lmlib/prompt/GemmaPromptStrategy.java
View File

@@ -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";
}
}

40
src/main/java/de/extio/lmlib/prompt/NoPromptStrategy.java Normal file
View File

@@ -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;
}
}

41
src/test/java/de/extio/lmlib/client/cached/CachedClientTest.java
View File

@@ -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(
ModelCategory.MEDIUM,
var completion = this.clientService.getClient(ModelCategory.SMALL).completion(
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(
ModelCategory.MEDIUM,
completion = this.clientService.getClient(ModelCategory.SMALL).completion(
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(
ModelCategory.MEDIUM,
completion = this.clientService.getClient(ModelCategory.SMALL).completion(
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());
}
}

43
src/test/java/de/extio/lmlib/client/oai/CompletionClientTest.java
View File

@@ -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;
@@ -40,7 +41,7 @@ public class CompletionClientTest {
@Autowired
private ChatCompletionClient chatCompletionclient;
@Test
void completion() throws Exception {
final var completion = this.textCompletionClient.completion(
@@ -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,13 +59,28 @@ 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
void chatCompletion() throws Exception {
final var conversation = Conversation.create("You are a helpful assistant", "I would like to paint my canvas in a solid color.");
@@ -74,9 +90,24 @@ 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
@Test
void hugePrompt() throws Exception {

8
src/test/resources/application-test.properties
View File

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

12
src/test/resources/local-chatcompletion.properties Executable file
View File

@@ -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

4
src/test/resources/local-completion.properties
View File

@@ -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=