Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add OpenAiChatModel stream observability

Browse Source 
Integrated Micrometer's Observation into the OpenAiChatModel#stream reactive chain.

Included changes:
 - Added ability to aggregate streaming responses for use in Observation metadata.
 - Improved error handling and logging for chat response processing.
 - Updated unit tests to include new observation logic and subscribe to Flux responses.
 - Refined validation of observations in both normal and streaming chat operations.
 - Disabled retry for streaming which used RetryTemplate - should use .retryWhen operator as the next step.
 - Added an integration test.

Resolves #1190

Co-authored-by Christian Tzolov <ctzolov@vmware.com>
This commit is contained in:
Dariusz Jędrzejczyk
2024-08-07 13:15:06 +02:00
committed by Christian Tzolov
parent 86348e4bb4
commit 478f180297
7 changed files with 246 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

117
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java
View File

@@ -39,6 +39,7 @@ import org.springframework.ai.chat.model.AbstractToolCallSupport;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.ai.chat.model.StreamingChatModel;
import org.springframework.ai.chat.observation.ChatModelObservationContext;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
@@ -72,7 +73,9 @@ import org.springframework.util.MimeType;
import org.springframework.util.MultiValueMap;
import org.springframework.util.StringUtils;
import io.micrometer.observation.Observation;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
@@ -271,64 +274,90 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode
@Override
public Flux<ChatResponse> stream(Prompt prompt) {
return Flux.deferContextual(contextView -> {
ChatCompletionRequest request = createRequest(prompt, true);
ChatCompletionRequest request = createRequest(prompt, true);
Flux<OpenAiApi.ChatCompletionChunk> completionChunks = this.openAiApi.chatCompletionStream(request,
getAdditionalHttpHeaders(prompt));
Flux<OpenAiApi.ChatCompletionChunk> completionChunks = this.retryTemplate
.execute(ctx -> this.openAiApi.chatCompletionStream(request, getAdditionalHttpHeaders(prompt)));
// For chunked responses, only the first chunk contains the choice role.
// The rest of the chunks with same ID share the same role.
ConcurrentHashMap<String, String> roleMap = new ConcurrentHashMap<>();
// For chunked responses, only the first chunk contains the choice role.
// The rest of the chunks with same ID share the same role.
ConcurrentHashMap<String, String> roleMap = new ConcurrentHashMap<>();
final ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
.prompt(prompt)
.operationMetadata(buildOperationMetadata())
.requestOptions(buildRequestOptions(request))
.build();
// Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse
// the function call handling logic.
Flux<ChatResponse> chatResponse = completionChunks.map(this::chunkToChatCompletion)
.switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
try {
@SuppressWarnings("null")
String id = chatCompletion2.id();
Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry);
// @formatter:off
List<Generation> generations = chatCompletion2.choices().stream().map(choice -> {
if (choice.message().role() != null) {
roleMap.putIfAbsent(id, choice.message().role().name());
}
Map<String, Object> metadata = Map.of(
"id", chatCompletion2.id(),
"role", roleMap.getOrDefault(id, ""),
"index", choice.index(),
"finishReason", choice.finishReason() != null ? choice.finishReason().name() : "");
return buildGeneration(choice, metadata);
observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
// Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse
// the function call handling logic.
Flux<ChatResponse> chatResponse = completionChunks.map(this::chunkToChatCompletion)
.switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
try {
@SuppressWarnings("null")
String id = chatCompletion2.id();
List<Generation> generations = chatCompletion2.choices().stream().map(choice -> {// @formatter:off
if (choice.message().role() != null) {
roleMap.putIfAbsent(id, choice.message().role().name());
}
Map<String, Object> metadata = Map.of(
"id", chatCompletion2.id(),
"role", roleMap.getOrDefault(id, ""),
"index", choice.index(),
"finishReason", choice.finishReason() != null ? choice.finishReason().name() : "");
return buildGeneration(choice, metadata);
}).toList();
// @formatter:on
// @formatter:on
if (chatCompletion2.usage() != null) {
return new ChatResponse(generations, from(chatCompletion2, null));
}
else {
return new ChatResponse(generations);
catch (Exception e) {
logger.error("Error processing chat completion", e);
return new ChatResponse(List.of());
}
}));
// @formatter:off
Flux<ChatResponse> flux = chatResponse.flatMap(response -> {
if (isToolCall(response, Set.of(OpenAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(),
OpenAiApi.ChatCompletionFinishReason.STOP.name()))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
// conversation that contains the call responses.
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
}
catch (Exception e) {
logger.error("Error processing chat completion", e);
return new ChatResponse(List.of());
else {
return Flux.just(response);
}
})
.doOnError(observation::error)
.doFinally(s -> {
// TODO: Consider a custom ObservationContext and
// include additional metadata
// if (s == SignalType.CANCEL) {
// observationContext.setAborted(true);
// }
observation.stop();
})
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
// @formatter:on
}));
return new MessageAggregator().aggregate(flux, mergedChatResponse -> {
observationContext.setResponse(mergedChatResponse);
});
return chatResponse.flatMap(response -> {
if (isToolCall(response, Set.of(OpenAiApi.ChatCompletionFinishReason.TOOL_CALLS.name(),
OpenAiApi.ChatCompletionFinishReason.STOP.name()))) {
var toolCallConversation = handleToolCalls(prompt, response);
// Recursively call the stream method with the tool call message
// conversation that contains the call responses.
return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
}
else {
return Flux.just(response);
}
});
}

8
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/MessageTypeContentTests.java
View File

@@ -104,7 +104,7 @@ public class MessageTypeContentTests {
when(openAiApi.chatCompletionStream(pomptCaptor.capture(), headersCaptor.capture())).thenReturn(fluxResponse);
chatModel.stream(new Prompt(List.of(new UserMessage("test message"))));
chatModel.stream(new Prompt(List.of(new UserMessage("test message")))).subscribe();
validateStringContent(pomptCaptor.getValue());
assertThat(headersCaptor.getValue()).isEmpty();
@@ -137,8 +137,10 @@ public class MessageTypeContentTests {
when(openAiApi.chatCompletionStream(pomptCaptor.capture(), headersCaptor.capture())).thenReturn(fluxResponse);
URL mediaUrl = new URL("http://test");
chatModel.stream(new Prompt(
List.of(new UserMessage("test message", List.of(new Media(MimeTypeUtils.IMAGE_JPEG, mediaUrl))))));
chatModel
.stream(new Prompt(
List.of(new UserMessage("test message", List.of(new Media(MimeTypeUtils.IMAGE_JPEG, mediaUrl))))))
.subscribe();
validateComplexContent(pomptCaptor.getValue());
}

51
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelObservationIT.java
View File

@@ -18,6 +18,9 @@ package org.springframework.ai.openai.chat;
import io.micrometer.common.KeyValue;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import reactor.core.publisher.Flux;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
@@ -37,6 +40,7 @@ import org.springframework.context.annotation.Bean;
import org.springframework.retry.support.RetryTemplate;
import java.util.List;
import java.util.stream.Collectors;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames;
@@ -57,8 +61,14 @@ public class OpenAiChatModelObservationIT {
@Autowired
OpenAiChatModel chatModel;
@BeforeEach
void beforeEach() {
observationRegistry.clear();
}
@Test
void observationForEmbeddingOperation() {
void observationForChatOperation() {
var options = OpenAiChatOptions.builder()
.withModel(OpenAiApi.ChatModel.GPT_4_O_MINI.getValue())
.withFrequencyPenalty(0f)
@@ -77,6 +87,45 @@ public class OpenAiChatModelObservationIT {
ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
assertThat(responseMetadata).isNotNull();
validate(responseMetadata);
}
@Test
void observationForStreamingChatOperation() {
var options = OpenAiChatOptions.builder()
.withModel(OpenAiApi.ChatModel.GPT_4_O_MINI.getValue())
.withFrequencyPenalty(0f)
.withMaxTokens(2048)
.withPresencePenalty(0f)
.withStop(List.of("this-is-the-end"))
.withTemperature(0.7f)
.withTopP(1f)
.withStreamUsage(true)
.build();
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
Flux<ChatResponse> chatResponseFlux = chatModel.stream(prompt);
List<ChatResponse> responses = chatResponseFlux.collectList().block();
assertThat(responses).isNotEmpty();
assertThat(responses).hasSizeGreaterThan(10);
String aggregatedResponse = responses.subList(0, responses.size() - 1)
.stream()
.map(r -> r.getResult().getOutput().getContent())
.collect(Collectors.joining());
assertThat(aggregatedResponse).isNotEmpty();
ChatResponse lastChatResponse = responses.get(responses.size() - 1);
ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata();
assertThat(responseMetadata).isNotNull();
validate(responseMetadata);
}
private void validate(ChatResponseMetadata responseMetadata) {
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()
.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)

3
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiPaymentTransactionIT.java
View File

@@ -53,8 +53,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
*/
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_PROJECT_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "VERTEX_AI_GEMINI_LOCATION", matches = ".*")
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
public class OpenAiPaymentTransactionIT {
private final static Logger logger = LoggerFactory.getLogger(OpenAiPaymentTransactionIT.class);

5
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiRetryTests.java
View File

@@ -19,6 +19,7 @@ import java.util.List;
import java.util.Optional;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
@@ -163,6 +164,7 @@ public class OpenAiRetryTests {
}
@Test
@Disabled("Currently stream() does not implmement retry")
public void openAiChatStreamTransientError() {
var choice = new ChatCompletionChunk.ChunkChoice(ChatCompletionFinishReason.STOP, 0,
@@ -184,10 +186,11 @@ public class OpenAiRetryTests {
}
@Test
@Disabled("Currently stream() does not implmement retry")
public void openAiChatStreamNonTransientError() {
when(openAiApi.chatCompletionStream(isA(ChatCompletionRequest.class), any()))
.thenThrow(new RuntimeException("Non Transient Error"));
assertThrows(RuntimeException.class, () -> chatModel.stream(new Prompt("text")));
assertThrows(RuntimeException.class, () -> chatModel.stream(new Prompt("text")).subscribe());
}
@Test

1
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMultipleFunctionCallsIT.java
View File

@@ -27,7 +27,6 @@ import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.DefaultChatClient;
import org.springframework.ai.openai.OpenAiTestConfiguration;
import org.springframework.ai.openai.api.tool.MockWeatherService;
import org.springframework.ai.openai.testutils.AbstractIT;

126
spring-ai-core/src/main/java/org/springframework/ai/chat/model/MessageAggregator.java
View File

@@ -24,6 +24,15 @@ import java.util.function.Consumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.EmptyRateLimit;
import org.springframework.ai.chat.metadata.PromptMetadata;
import org.springframework.ai.chat.metadata.RateLimit;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
/**
@@ -40,32 +49,123 @@ public class MessageAggregator {
public Flux<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
Consumer<ChatResponse> onAggregationComplete) {
AtomicReference<StringBuilder> stringBufferRef = new AtomicReference<>(new StringBuilder());
AtomicReference<Map<String, Object>> mapRef = new AtomicReference<>();
// Assistant Message
AtomicReference<StringBuilder> messageTextContentRef = new AtomicReference<>(new StringBuilder());
AtomicReference<Map<String, Object>> messageMetadataMapRef = new AtomicReference<>();
// ChatGeneration Metadata
AtomicReference<ChatGenerationMetadata> generationMetadataRef = new AtomicReference<>(
ChatGenerationMetadata.NULL);
// Usage
AtomicReference<Long> metadataUsagePromptTokensRef = new AtomicReference<>(0L);
AtomicReference<Long> metadataUsageGenerationTokensRef = new AtomicReference<>(0L);
AtomicReference<Long> metadataUsageTotalTokensRef = new AtomicReference<>(0L);
AtomicReference<PromptMetadata> metadataPromptMetadataRef = new AtomicReference<>(PromptMetadata.empty());
AtomicReference<RateLimit> metadataRateLimitRef = new AtomicReference<>(new EmptyRateLimit());
AtomicReference<String> metadataIdRef = new AtomicReference<>("");
AtomicReference<String> metadataModelRef = new AtomicReference<>("");
return fluxChatResponse.doOnSubscribe(subscription -> {
// logger.info("Aggregation Subscribe:" + subscription);
stringBufferRef.set(new StringBuilder());
mapRef.set(new HashMap<>());
messageTextContentRef.set(new StringBuilder());
messageMetadataMapRef.set(new HashMap<>());
metadataIdRef.set("");
metadataModelRef.set("");
metadataUsagePromptTokensRef.set(0L);
metadataUsageGenerationTokensRef.set(0L);
metadataUsageTotalTokensRef.set(0L);
metadataPromptMetadataRef.set(PromptMetadata.empty());
metadataRateLimitRef.set(new EmptyRateLimit());
}).doOnNext(chatResponse -> {
// logger.info("Aggregation Next:" + chatResponse);
if (chatResponse.getResult() != null) {
if (chatResponse.getResult().getMetadata() != null
&& chatResponse.getResult().getMetadata() != ChatGenerationMetadata.NULL) {
generationMetadataRef.set(chatResponse.getResult().getMetadata());
}
if (chatResponse.getResult().getOutput().getContent() != null) {
stringBufferRef.get().append(chatResponse.getResult().getOutput().getContent());
messageTextContentRef.get().append(chatResponse.getResult().getOutput().getContent());
}
if (chatResponse.getResult().getOutput().getMetadata() != null) {
mapRef.get().putAll(chatResponse.getResult().getOutput().getMetadata());
messageMetadataMapRef.get().putAll(chatResponse.getResult().getOutput().getMetadata());
}
}
if (chatResponse.getMetadata() != null) {
if (chatResponse.getMetadata().getUsage() != null) {
Usage usage = chatResponse.getMetadata().getUsage();
metadataUsagePromptTokensRef.set(
usage.getPromptTokens() > 0 ? usage.getPromptTokens() : metadataUsagePromptTokensRef.get());
metadataUsageGenerationTokensRef.set(usage.getGenerationTokens() > 0 ? usage.getGenerationTokens()
: metadataUsageGenerationTokensRef.get());
metadataUsageTotalTokensRef
.set(usage.getTotalTokens() > 0 ? usage.getTotalTokens() : metadataUsageTotalTokensRef.get());
}
if (chatResponse.getMetadata().getPromptMetadata() != null
&& chatResponse.getMetadata().getPromptMetadata().iterator().hasNext()) {
metadataPromptMetadataRef.set(chatResponse.getMetadata().getPromptMetadata());
}
if (chatResponse.getMetadata().getRateLimit() != null
&& !(metadataRateLimitRef.get() instanceof EmptyRateLimit)) {
metadataRateLimitRef.set(chatResponse.getMetadata().getRateLimit());
}
if (StringUtils.hasText(chatResponse.getMetadata().getId())) {
metadataIdRef.set(chatResponse.getMetadata().getId());
}
if (StringUtils.hasText(chatResponse.getMetadata().getModel())) {
metadataModelRef.set(chatResponse.getMetadata().getModel());
}
}
}).doOnComplete(() -> {
// logger.debug("Aggregation Complete");
onAggregationComplete
.accept(new ChatResponse(List.of(new Generation(stringBufferRef.get().toString(), mapRef.get()))));
stringBufferRef.set(new StringBuilder());
mapRef.set(new HashMap<>());
var usage = new DefaultUsage(metadataUsagePromptTokensRef.get(), metadataUsageGenerationTokensRef.get(),
metadataUsageTotalTokensRef.get());
var chatResponseMetadata = ChatResponseMetadata.builder()
.withId(metadataIdRef.get())
.withModel(metadataModelRef.get())
.withRateLimit(metadataRateLimitRef.get())
.withUsage(usage)
.withPromptMetadata(metadataPromptMetadataRef.get())
.build();
onAggregationComplete.accept(new ChatResponse(List.of(new Generation(
new AssistantMessage(messageTextContentRef.get().toString(), messageMetadataMapRef.get()),
generationMetadataRef.get())), chatResponseMetadata));
messageTextContentRef.set(new StringBuilder());
messageMetadataMapRef.set(new HashMap<>());
metadataIdRef.set("");
metadataModelRef.set("");
metadataUsagePromptTokensRef.set(0L);
metadataUsageGenerationTokensRef.set(0L);
metadataUsageTotalTokensRef.set(0L);
metadataPromptMetadataRef.set(PromptMetadata.empty());
metadataRateLimitRef.set(new EmptyRateLimit());
}).doOnError(e -> {
logger.error("Aggregation Error", e);
});
}
public record DefaultUsage(long promptTokens, long generationTokens, long totalTokens) implements Usage {
@Override
public Long getPromptTokens() {
return promptTokens();
}
@Override
public Long getGenerationTokens() {
return generationTokens();
}
@Override
public Long getTotalTokens() {
return totalTokens();
}
}
}