From 5e7cc654a3f324e9378d3dab31451737925f4e44 Mon Sep 17 00:00:00 2001 From: Jonatan Ivanov Date: Thu, 18 Sep 2025 13:10:53 -0700 Subject: [PATCH 1/2] Log chat client completion data and make the response text available in the Observation Context --- .../ChatClientAutoConfiguration.java | 26 ++++ .../ChatClientBuilderProperties.java | 23 ++++ ...ientObservationAutoConfigurationTests.java | 128 ++++++++++++++++-- spring-ai-client-chat/pom.xml | 6 + .../ai/chat/client/DefaultChatClient.java | 6 +- ...hatClientCompletionObservationHandler.java | 61 +++++++++ .../ChatClientObservationContext.java | 40 ++++++ .../chat/client/DefaultChatClientTests.java | 69 ++++++++++ ...ientCompletionObservationHandlerTests.java | 109 +++++++++++++++ .../ChatClientObservationContextTests.java | 74 ++++++++++ .../ROOT/pages/observability/index.adoc | 9 +- 11 files changed, 538 insertions(+), 13 deletions(-) create mode 100644 spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java create mode 100644 spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandlerTests.java diff --git a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientAutoConfiguration.java b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientAutoConfiguration.java index 32a09a85425..de688a542d7 100644 --- a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientAutoConfiguration.java +++ b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientAutoConfiguration.java @@ -23,6 +23,7 @@ import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.client.ChatClientCustomizer; +import org.springframework.ai.chat.client.observation.ChatClientCompletionObservationHandler; import org.springframework.ai.chat.client.observation.ChatClientObservationContext; import org.springframework.ai.chat.client.observation.ChatClientObservationConvention; import org.springframework.ai.chat.client.observation.ChatClientPromptContentObservationHandler; @@ -71,6 +72,11 @@ private static void logPromptContentWarning() { "You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!"); } + private static void logCompletionWarning() { + logger.warn( + "You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!"); + } + @Bean @ConditionalOnMissingBean ChatClientBuilderConfigurer chatClientBuilderConfigurer(ObjectProvider customizerProvider) { @@ -108,6 +114,17 @@ TracingAwareLoggingObservationHandler chatClientPr return new TracingAwareLoggingObservationHandler<>(new ChatClientPromptContentObservationHandler(), tracer); } + @Bean + @ConditionalOnMissingBean(value = ChatClientCompletionObservationHandler.class, + name = "chatClientCompletionObservationHandler") + @ConditionalOnProperty(prefix = ChatClientBuilderProperties.CONFIG_PREFIX + ".observations", + name = "log-completion", havingValue = "true") + TracingAwareLoggingObservationHandler chatClientCompletionObservationHandler( + Tracer tracer) { + logCompletionWarning(); + return new TracingAwareLoggingObservationHandler<>(new ChatClientCompletionObservationHandler(), tracer); + } + } @Configuration(proxyBeanMethods = false) @@ -123,6 +140,15 @@ ChatClientPromptContentObservationHandler chatClientPromptContentObservationHand return new ChatClientPromptContentObservationHandler(); } + @Bean + @ConditionalOnMissingBean + @ConditionalOnProperty(prefix = ChatClientBuilderProperties.CONFIG_PREFIX + ".observations", + name = "log-completion", havingValue = "true") + ChatClientCompletionObservationHandler chatClientCompletionObservationHandler() { + logCompletionWarning(); + return new ChatClientCompletionObservationHandler(); + } + } } diff --git a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientBuilderProperties.java b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientBuilderProperties.java index 9e58092390a..f2e5b09c277 100644 --- a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientBuilderProperties.java +++ b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientBuilderProperties.java @@ -26,6 +26,7 @@ * @author Josh Long * @author Arjen Poutsma * @author Thomas Vitale + * @author Jonatan Ivanov * @since 1.0.0 */ @ConfigurationProperties(ChatClientBuilderProperties.CONFIG_PREFIX) @@ -59,14 +60,36 @@ public static class Observations { */ private boolean logPrompt = false; + /** + * Whether to log the completion content in the observations. + * @since 1.1.0 + */ + private boolean logCompletion = false; + public boolean isLogPrompt() { return this.logPrompt; } + /** + * @return Whether logging completion data is enabled or not. + * @since 1.1.0 + */ + public boolean isLogCompletion() { + return this.logCompletion; + } + public void setLogPrompt(boolean logPrompt) { this.logPrompt = logPrompt; } + /** + * @param logCompletion should completion data logging be enabled or not. + * @since 1.1.0 + */ + public void setLogCompletion(boolean logCompletion) { + this.logCompletion = logCompletion; + } + } } diff --git a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/test/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientObservationAutoConfigurationTests.java b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/test/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientObservationAutoConfigurationTests.java index 94b18d1fa83..4f076b4ef64 100644 --- a/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/test/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientObservationAutoConfigurationTests.java +++ b/auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/test/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientObservationAutoConfigurationTests.java @@ -20,6 +20,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.springframework.ai.chat.client.observation.ChatClientCompletionObservationHandler; import org.springframework.ai.chat.client.observation.ChatClientObservationContext; import org.springframework.ai.chat.client.observation.ChatClientPromptContentObservationHandler; import org.springframework.ai.observation.TracingAwareLoggingObservationHandler; @@ -48,16 +49,18 @@ class ChatClientObservationAutoConfigurationTests { .withConfiguration(AutoConfigurations.of(ChatClientAutoConfiguration.class)); @Test - void promptContentHandlerNoTracer() { + void handlersNoTracer() { this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @Test - void promptContentHandlerWithTracer() { + void handlersWithTracer() { this.contextRunner.withUserConfiguration(TracerConfiguration.class) .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @@ -66,6 +69,7 @@ void promptContentHandlerEnabledNoTracer(CapturedOutput output) { this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) .withPropertyValues("spring.ai.chat.client.observations.log-prompt=true") .run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); assertThat(output).contains( "You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!"); @@ -76,6 +80,7 @@ void promptContentHandlerEnabledWithTracer(CapturedOutput output) { this.contextRunner.withUserConfiguration(TracerConfiguration.class) .withPropertyValues("spring.ai.chat.client.observations.log-prompt=true") .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .hasSingleBean(TracingAwareLoggingObservationHandler.class)); assertThat(output).contains( "You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!"); @@ -86,6 +91,7 @@ void promptContentHandlerDisabledNoTracer() { this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) .withPropertyValues("spring.ai.chat.client.observations.log-prompt=false") .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @@ -94,6 +100,47 @@ void promptContentHandlerDisabledWithTracer() { this.contextRunner.withUserConfiguration(TracerConfiguration.class) .withPropertyValues("spring.ai.chat.client.observations.log-prompt=false") .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) + .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); + } + + @Test + void completionHandlerEnabledNoTracer(CapturedOutput output) { + this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=true") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .hasSingleBean(ChatClientCompletionObservationHandler.class) + .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); + assertThat(output).contains( + "You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!"); + } + + @Test + void completionHandlerEnabledWithTracer(CapturedOutput output) { + this.contextRunner.withUserConfiguration(TracerConfiguration.class) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=true") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) + .hasSingleBean(TracingAwareLoggingObservationHandler.class)); + assertThat(output).contains( + "You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!"); + } + + @Test + void completionHandlerDisabledNoTracer() { + this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=false") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) + .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); + } + + @Test + void completionDisabledWithTracer() { + this.contextRunner.withUserConfiguration(TracerConfiguration.class) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=false") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @@ -104,6 +151,7 @@ void customChatClientPromptContentObservationHandlerNoTracer() { .withPropertyValues("spring.ai.chat.client.observations.log-prompt=true") .run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class) .hasBean("customChatClientPromptContentObservationHandler") + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @@ -114,20 +162,61 @@ void customChatClientPromptContentObservationHandlerWithTracer() { .withPropertyValues("spring.ai.chat.client.observations.log-prompt=true") .run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class) .hasBean("customChatClientPromptContentObservationHandler") + .doesNotHaveBean(ChatClientCompletionObservationHandler.class) .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); } @Test - void customTracingAwareLoggingObservationHandler() { + void customTracingAwareLoggingObservationHandlerForChatClientPromptContent() { this.contextRunner.withUserConfiguration(TracerConfiguration.class) - .withUserConfiguration(CustomTracingAwareLoggingObservationHandlerConfiguration.class) + .withUserConfiguration( + CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration.class) .withPropertyValues("spring.ai.chat.client.observations.log-prompt=true") .run(context -> { assertThat(context).hasSingleBean(TracingAwareLoggingObservationHandler.class) .hasBean("chatClientPromptContentObservationHandler") - .doesNotHaveBean(ChatClientPromptContentObservationHandler.class); - assertThat(context.getBean(TracingAwareLoggingObservationHandler.class)) - .isSameAs(CustomTracingAwareLoggingObservationHandlerConfiguration.handlerInstance); + .doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class); + assertThat(context.getBean(TracingAwareLoggingObservationHandler.class)).isSameAs( + CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration.handlerInstance); + }); + } + + @Test + void customChatClientCompletionObservationHandlerNoTracer() { + this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class)) + .withUserConfiguration(CustomChatClientCompletionObservationHandlerConfiguration.class) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=true") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .hasSingleBean(ChatClientCompletionObservationHandler.class) + .hasBean("customChatClientCompletionObservationHandler") + .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); + } + + @Test + void customChatClientCompletionObservationHandlerWithTracer() { + this.contextRunner.withUserConfiguration(TracerConfiguration.class) + .withUserConfiguration(CustomChatClientCompletionObservationHandlerConfiguration.class) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=true") + .run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .hasSingleBean(ChatClientCompletionObservationHandler.class) + .hasBean("customChatClientCompletionObservationHandler") + .doesNotHaveBean(TracingAwareLoggingObservationHandler.class)); + } + + @Test + void customTracingAwareLoggingObservationHandlerForChatClientCompletion() { + this.contextRunner.withUserConfiguration(TracerConfiguration.class) + .withUserConfiguration( + CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration.class) + .withPropertyValues("spring.ai.chat.client.observations.log-completion=true") + .run(context -> { + assertThat(context).hasSingleBean(TracingAwareLoggingObservationHandler.class) + .hasBean("chatClientCompletionObservationHandler") + .doesNotHaveBean(ChatClientPromptContentObservationHandler.class) + .doesNotHaveBean(ChatClientCompletionObservationHandler.class); + assertThat(context.getBean(TracingAwareLoggingObservationHandler.class)).isSameAs( + CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration.handlerInstance); }); } @@ -152,7 +241,7 @@ ChatClientPromptContentObservationHandler customChatClientPromptContentObservati } @Configuration(proxyBeanMethods = false) - static class CustomTracingAwareLoggingObservationHandlerConfiguration { + static class CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration { static TracingAwareLoggingObservationHandler handlerInstance = new TracingAwareLoggingObservationHandler<>( new ChatClientPromptContentObservationHandler(), null); @@ -164,4 +253,27 @@ TracingAwareLoggingObservationHandler chatClientPr } + @Configuration(proxyBeanMethods = false) + static class CustomChatClientCompletionObservationHandlerConfiguration { + + @Bean + ChatClientCompletionObservationHandler customChatClientCompletionObservationHandler() { + return new ChatClientCompletionObservationHandler(); + } + + } + + @Configuration(proxyBeanMethods = false) + static class CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration { + + static TracingAwareLoggingObservationHandler handlerInstance = new TracingAwareLoggingObservationHandler<>( + new ChatClientCompletionObservationHandler(), null); + + @Bean + TracingAwareLoggingObservationHandler chatClientCompletionObservationHandler() { + return handlerInstance; + } + + } + } diff --git a/spring-ai-client-chat/pom.xml b/spring-ai-client-chat/pom.xml index 5253a775a01..5b4242ef5da 100644 --- a/spring-ai-client-chat/pom.xml +++ b/spring-ai-client-chat/pom.xml @@ -104,6 +104,12 @@ test + + io.micrometer + micrometer-observation-test + test + + com.fasterxml.jackson.module jackson-module-kotlin diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java index 20b207d5c5c..7f0ecd2e00d 100644 --- a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java @@ -76,6 +76,7 @@ * @author Soby Chacko * @author Dariusz Jedrzejczyk * @author Thomas Vitale + * @author Jonatan Ivanov * @since 1.0.0 */ public class DefaultChatClient implements ChatClient { @@ -513,7 +514,9 @@ private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest c // CHECKSTYLE:OFF var chatClientResponse = observation.observe(() -> { // Apply the advisor chain that terminates with the ChatModelCallAdvisor. - return this.advisorChain.nextCall(chatClientRequest); + var response = this.advisorChain.nextCall(chatClientRequest); + observationContext.setResponse(response); + return response; }); // CHECKSTYLE:ON return chatClientResponse != null ? chatClientResponse : ChatClientResponse.builder().build(); @@ -572,6 +575,7 @@ private Flux doGetObservableFluxChatResponse(ChatClientReque // @formatter:off // Apply the advisor chain that terminates with the ChatModelStreamAdvisor. return this.advisorChain.nextStream(chatClientRequest) + .doOnNext(observationContext::appendResponse) .doOnError(observation::error) .doFinally(s -> observation.stop()) .contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation)); diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java new file mode 100644 index 00000000000..44684b0d36d --- /dev/null +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java @@ -0,0 +1,61 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.client.observation; + +import java.util.Collections; +import java.util.List; + +import io.micrometer.observation.Observation; +import io.micrometer.observation.ObservationHandler; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.observation.ObservabilityHelper; +import org.springframework.util.StringUtils; + +/** + * Handler for emitting the chat client completion content to logs. + * + * @author Jonatan Ivanov + * @since 1.1.0 + */ +public class ChatClientCompletionObservationHandler implements ObservationHandler { + + private static final Logger logger = LoggerFactory.getLogger(ChatClientCompletionObservationHandler.class); + + @Override + public void onStop(ChatClientObservationContext context) { + logger.info("Chat Client Completion:\n{}", ObservabilityHelper.concatenateStrings(completion(context))); + } + + private List completion(ChatClientObservationContext context) { + if (context.getResponseText() == null) { + return List.of(); + } + else if (!StringUtils.hasText(context.getResponseText())) { + return List.of(); + } + + return Collections.singletonList(context.getResponseText()); + } + + @Override + public boolean supportsContext(Observation.Context context) { + return context instanceof ChatClientObservationContext; + } + +} diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java index df750398a0f..46102098083 100644 --- a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java @@ -17,12 +17,16 @@ package org.springframework.ai.chat.client.observation; import java.util.List; +import java.util.stream.Collectors; import io.micrometer.observation.Observation; import org.springframework.ai.chat.client.ChatClientAttributes; import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.Advisor; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.model.Generation; import org.springframework.ai.observation.AiOperationMetadata; import org.springframework.ai.observation.conventions.AiOperationType; import org.springframework.ai.observation.conventions.AiProvider; @@ -35,12 +39,16 @@ * * @author Christian Tzolov * @author Thomas Vitale + * @author Jonatan Ivanov * @since 1.0.0 */ public class ChatClientObservationContext extends Observation.Context { private final ChatClientRequest request; + @Nullable + private String responseText; + private final AiOperationMetadata operationMetadata = new AiOperationMetadata(AiOperationType.FRAMEWORK.value(), AiProvider.SPRING_AI.value()); @@ -86,6 +94,38 @@ public String getFormat() { return null; } + @Nullable + String getResponseText() { + return this.responseText; + } + + /** + * @param response Chat client response to record. + * @since 1.1.0 + */ + public void setResponse(ChatClientResponse response) { + if (response.chatResponse() != null) { + this.responseText = generationsToString(response.chatResponse().getResults()); + } + } + + /** + * @param response Chat client response to record. + * @since 1.1.0 + */ + public void appendResponse(ChatClientResponse response) { + if (this.responseText == null) { + setResponse(response); + } + else if (response.chatResponse() != null) { + this.responseText += generationsToString(response.chatResponse().getResults()); + } + } + + private String generationsToString(List generations) { + return generations.stream().map(Generation::getOutput).map(Message::getText).collect(Collectors.joining("\n")); + } + public static final class Builder { private ChatClientRequest chatClientRequest; diff --git a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java index 07adcf72b48..17a803be7fe 100644 --- a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java +++ b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java @@ -27,6 +27,7 @@ import java.util.function.Consumer; import io.micrometer.observation.ObservationRegistry; +import io.micrometer.observation.tck.TestObservationRegistry; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -35,6 +36,7 @@ import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor; import org.springframework.ai.chat.client.advisor.api.Advisor; import org.springframework.ai.chat.client.advisor.api.BaseAdvisorChain; +import org.springframework.ai.chat.client.observation.ChatClientObservationContext; import org.springframework.ai.chat.client.observation.ChatClientObservationConvention; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; @@ -67,6 +69,7 @@ * Unit tests for {@link DefaultChatClient}. * * @author Thomas Vitale + * @author Jonatan Ivanov */ class DefaultChatClientTests { @@ -837,6 +840,39 @@ void whenSimplePromptThenChatClientResponse() { assertThat(actualPrompt.getInstructions().get(0).getText()).isEqualTo("my question"); } + @Test + void whenSimplePromptThenSetRequestAndResponseOnObservationContext() { + ChatModel chatModel = mock(ChatModel.class); + TestObservationRegistry observationRegistry = TestObservationRegistry.create(); + ArgumentCaptor promptCaptor = ArgumentCaptor.forClass(Prompt.class); + given(chatModel.call(promptCaptor.capture())) + .willReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("response"))))); + + ChatClient chatClient = new DefaultChatClientBuilder(chatModel, observationRegistry, null).build(); + DefaultChatClient.DefaultChatClientRequestSpec chatClientRequestSpec = (DefaultChatClient.DefaultChatClientRequestSpec) chatClient + .prompt("my question"); + DefaultChatClient.DefaultCallResponseSpec spec = (DefaultChatClient.DefaultCallResponseSpec) chatClientRequestSpec + .call(); + + ChatClientResponse chatClientResponse = spec.chatClientResponse(); + assertThat(chatClientResponse).isNotNull(); + + ChatResponse chatResponse = chatClientResponse.chatResponse(); + assertThat(chatResponse).isNotNull(); + assertThat(chatResponse.getResult().getOutput().getText()).isEqualTo("response"); + + Prompt actualPrompt = promptCaptor.getValue(); + assertThat(actualPrompt.getInstructions()).hasSize(1); + assertThat(actualPrompt.getInstructions().get(0).getText()).isEqualTo("my question"); + + assertThat(observationRegistry).hasObservationWithNameEqualTo("spring.ai.chat.client") + .that() + .isInstanceOfSatisfying(ChatClientObservationContext.class, context -> { + assertThat(context.getRequest().prompt()).isEqualTo(actualPrompt); + assertThat(context).extracting("responseText").isEqualTo("response"); + }); + } + @Test void whenSimplePromptThenChatResponse() { ChatModel chatModel = mock(ChatModel.class); @@ -1350,6 +1386,39 @@ void whenSimplePromptThenFluxChatClientResponse() { assertThat(actualPrompt.getInstructions().get(0).getText()).isEqualTo("my question"); } + @Test + void whenSimplePromptThenSetFluxResponseOnObservationContext() { + ChatModel chatModel = mock(ChatModel.class); + TestObservationRegistry observationRegistry = TestObservationRegistry.create(); + ArgumentCaptor promptCaptor = ArgumentCaptor.forClass(Prompt.class); + given(chatModel.stream(promptCaptor.capture())) + .willReturn(Flux.just(new ChatResponse(List.of(new Generation(new AssistantMessage("response")))))); + + ChatClient chatClient = new DefaultChatClientBuilder(chatModel, observationRegistry, null).build(); + DefaultChatClient.DefaultChatClientRequestSpec chatClientRequestSpec = (DefaultChatClient.DefaultChatClientRequestSpec) chatClient + .prompt("my question"); + DefaultChatClient.DefaultStreamResponseSpec spec = (DefaultChatClient.DefaultStreamResponseSpec) chatClientRequestSpec + .stream(); + + ChatClientResponse chatClientResponse = spec.chatClientResponse().blockLast(); + assertThat(chatClientResponse).isNotNull(); + + ChatResponse chatResponse = chatClientResponse.chatResponse(); + assertThat(chatResponse).isNotNull(); + assertThat(chatResponse.getResult().getOutput().getText()).isEqualTo("response"); + + Prompt actualPrompt = promptCaptor.getValue(); + assertThat(actualPrompt.getInstructions()).hasSize(1); + assertThat(actualPrompt.getInstructions().get(0).getText()).isEqualTo("my question"); + + assertThat(observationRegistry).hasObservationWithNameEqualTo("spring.ai.chat.client") + .that() + .isInstanceOfSatisfying(ChatClientObservationContext.class, context -> { + assertThat(context.getRequest().prompt()).isEqualTo(actualPrompt); + assertThat(context).extracting("responseText").isEqualTo("response"); + }); + } + @Test void whenSimplePromptThenFluxChatResponse() { ChatModel chatModel = mock(ChatModel.class); diff --git a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandlerTests.java b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandlerTests.java new file mode 100644 index 00000000000..eb93dc1bfc2 --- /dev/null +++ b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandlerTests.java @@ -0,0 +1,109 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.client.observation; + +import java.util.List; + +import io.micrometer.observation.Observation; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; + +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.boot.test.system.CapturedOutput; +import org.springframework.boot.test.system.OutputCaptureExtension; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link ChatClientCompletionObservationHandler}. + * + * @author Jonatan Ivanov + */ +@ExtendWith(OutputCaptureExtension.class) +class ChatClientCompletionObservationHandlerTests { + + private final ChatClientCompletionObservationHandler observationHandler = new ChatClientCompletionObservationHandler(); + + @Test + void whenNotSupportedObservationContextThenReturnFalse() { + var context = new Observation.Context(); + assertThat(this.observationHandler.supportsContext(context)).isFalse(); + } + + @Test + void whenSupportedObservationContextThenReturnTrue() { + var context = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt(List.of())).build()) + .build(); + assertThat(this.observationHandler.supportsContext(context)).isTrue(); + } + + @Test + void whenEmptyResponseThenOutputNothing(CapturedOutput output) { + var context = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt(List.of())).build()) + .build(); + var response = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder().generations(List.of(new Generation(new AssistantMessage("")))).build()) + .build(); + context.setResponse(response); + + this.observationHandler.onStop(context); + assertThat(output).contains(""" + INFO o.s.a.c.c.o.ChatClientCompletionObservationHandler -- Chat Client Completion: + [] + """); + } + + @Test + void whenNullResponseThenOutputNothing(CapturedOutput output) { + var context = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt(List.of())).build()) + .build(); + + this.observationHandler.onStop(context); + assertThat(output).contains(""" + INFO o.s.a.c.c.o.ChatClientCompletionObservationHandler -- Chat Client Completion: + [] + """); + } + + @Test + void whenResponseWithTextThenOutputIt(CapturedOutput output) { + var context = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt(List.of())).build()) + .build(); + var response = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage("Test message")))) + .build()) + .build(); + context.setResponse(response); + + this.observationHandler.onStop(context); + assertThat(output).contains(""" + INFO o.s.a.c.c.o.ChatClientCompletionObservationHandler -- Chat Client Completion: + ["Test message"] + """); + } + +} diff --git a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java index f0361f803cb..0a63d392ece 100644 --- a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java +++ b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java @@ -25,8 +25,12 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.Advisor; +import org.springframework.ai.chat.messages.AssistantMessage; 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.prompt.Prompt; import static org.assertj.core.api.Assertions.assertThat; @@ -38,6 +42,7 @@ * * @author Christian Tzolov * @author Thomas Vitale + * @author Jonatan Ivanov */ @ExtendWith(MockitoExtension.class) class ChatClientObservationContextTests { @@ -184,4 +189,73 @@ void whenNoAdvisorsSpecifiedThenGetAdvisorsReturnsEmptyOrNull() { advisors -> assertThat(advisors).isEmpty()); } + @Test + void whenSetChatClientResponseThenReturnResponseText() { + var observationContext = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) + .build(); + var response = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage("Test message")))) + .build()) + .build(); + + observationContext.setResponse(response); + assertThat(observationContext.getResponseText()).isEqualTo("Test message"); + } + + @Test + void whenSetChatClientResponseWithNullChatResponseThenReturnNull() { + var observationContext = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) + .build(); + var nullResponse = ChatClientResponse.builder().chatResponse(null).build(); + + observationContext.setResponse(nullResponse); + assertThat(observationContext.getResponseText()).isNull(); + } + + @Test + void whenAppendChatClientResponseThenReturnResponseText() { + var observationContext = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) + .build(); + var response1 = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage("Test message ")))) + .build()) + .build(); + var response2 = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage("and another test message")))) + .build()) + .build(); + + observationContext.appendResponse(response1); + assertThat(observationContext.getResponseText()).isEqualTo("Test message "); + + observationContext.appendResponse(response2); + assertThat(observationContext.getResponseText()).isEqualTo("Test message and another test message"); + } + + @Test + void whenAppendChatClientResponseWithNullChatResponseThenReturnNull() { + var observationContext = ChatClientObservationContext.builder() + .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) + .build(); + var nullResponse = ChatClientResponse.builder().chatResponse(null).build(); + var response = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List.of(new Generation(new AssistantMessage("Test message ")))) + .build()) + .build(); + + observationContext.appendResponse(nullResponse); + assertThat(observationContext.getResponseText()).isNull(); + observationContext.appendResponse(response); + assertThat(observationContext.getResponseText()).isEqualTo("Test message "); + observationContext.appendResponse(nullResponse); + assertThat(observationContext.getResponseText()).isEqualTo("Test message "); + } + } diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/observability/index.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/observability/index.adoc index 2ddeacadd46..8bd76d8d851 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/observability/index.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/observability/index.adoc @@ -54,21 +54,22 @@ They measure the time spent performing the invocation and propagate the related |`spring.ai.chat.client.user.text` (deprecated) | Chat client user text. Optional. Superseded by `gen_ai.prompt`. |=== -=== Prompt Content +=== Prompt and Completion Data -The `ChatClient` prompt content is typically big and possibly containing sensitive information. +The `ChatClient` prompt and completion data is typically big and possibly containing sensitive information. For those reasons, it is not exported by default. -Spring AI supports logging the prompt content to help with debugging and troubleshooting. +Spring AI supports logging the prompt and completion data to help with debugging and troubleshooting. [cols="6,3,1", stripes=even] |==== | Property | Description | Default | `spring.ai.chat.client.observations.log-prompt` | Whether to log the chat client prompt content. | `false` +| `spring.ai.chat.client.observations.log-completion` | Whether to log the chat client completion content. | `false` |==== -WARNING: If you enable logging of the chat client prompt content, there's a risk of exposing sensitive or private information. Please, be careful! +WARNING: If you enable logging of the chat client prompt and completion data, there's a risk of exposing sensitive or private information. Please, be careful! === Input Data (Deprecated) From 09bbafc542c35f81bf298e19cfa742af3f554c86 Mon Sep 17 00:00:00 2001 From: Jonatan Ivanov Date: Thu, 25 Sep 2025 14:33:23 -0700 Subject: [PATCH 2/2] Utilize ChatClientMessageAggregator to set the response in the observation context --- .../ai/chat/client/DefaultChatClient.java | 7 ++- ...hatClientCompletionObservationHandler.java | 17 +++--- .../ChatClientObservationContext.java | 32 +++--------- .../chat/client/DefaultChatClientTests.java | 5 +- .../ChatClientObservationContextTests.java | 52 ++----------------- 5 files changed, 30 insertions(+), 83 deletions(-) diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java index 7f0ecd2e00d..64d4144b162 100644 --- a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java @@ -85,6 +85,8 @@ public class DefaultChatClient implements ChatClient { private static final TemplateRenderer DEFAULT_TEMPLATE_RENDERER = StTemplateRenderer.builder().build(); + private static final ChatClientMessageAggregator CHAT_CLIENT_MESSAGE_AGGREGATOR = new ChatClientMessageAggregator(); + private final DefaultChatClientRequestSpec defaultChatClientRequest; public DefaultChatClient(DefaultChatClientRequestSpec defaultChatClientRequest) { @@ -574,12 +576,13 @@ private Flux doGetObservableFluxChatResponse(ChatClientReque // @formatter:off // Apply the advisor chain that terminates with the ChatModelStreamAdvisor. - return this.advisorChain.nextStream(chatClientRequest) - .doOnNext(observationContext::appendResponse) + Flux chatClientResponse = this.advisorChain.nextStream(chatClientRequest) .doOnError(observation::error) .doFinally(s -> observation.stop()) .contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation)); // @formatter:on + return CHAT_CLIENT_MESSAGE_AGGREGATOR.aggregateChatClientResponse(chatClientResponse, + observationContext::setResponse); }); } diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java index 44684b0d36d..a42e2c8abe0 100644 --- a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientCompletionObservationHandler.java @@ -16,7 +16,6 @@ package org.springframework.ai.chat.client.observation; -import java.util.Collections; import java.util.List; import io.micrometer.observation.Observation; @@ -24,6 +23,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.model.Generation; import org.springframework.ai.observation.ObservabilityHelper; import org.springframework.util.StringUtils; @@ -43,14 +44,18 @@ public void onStop(ChatClientObservationContext context) { } private List completion(ChatClientObservationContext context) { - if (context.getResponseText() == null) { - return List.of(); - } - else if (!StringUtils.hasText(context.getResponseText())) { + if (context.getResponse() == null || context.getResponse().chatResponse() == null) { return List.of(); } - return Collections.singletonList(context.getResponseText()); + return context.getResponse() + .chatResponse() + .getResults() + .stream() + .map(Generation::getOutput) + .map(Message::getText) + .filter(StringUtils::hasText) + .toList(); } @Override diff --git a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java index 46102098083..c056904df58 100644 --- a/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java +++ b/spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/observation/ChatClientObservationContext.java @@ -17,7 +17,6 @@ package org.springframework.ai.chat.client.observation; import java.util.List; -import java.util.stream.Collectors; import io.micrometer.observation.Observation; @@ -25,8 +24,6 @@ import org.springframework.ai.chat.client.ChatClientRequest; import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.client.advisor.api.Advisor; -import org.springframework.ai.chat.messages.Message; -import org.springframework.ai.chat.model.Generation; import org.springframework.ai.observation.AiOperationMetadata; import org.springframework.ai.observation.conventions.AiOperationType; import org.springframework.ai.observation.conventions.AiProvider; @@ -47,7 +44,7 @@ public class ChatClientObservationContext extends Observation.Context { private final ChatClientRequest request; @Nullable - private String responseText; + private ChatClientResponse response; private final AiOperationMetadata operationMetadata = new AiOperationMetadata(AiOperationType.FRAMEWORK.value(), AiProvider.SPRING_AI.value()); @@ -94,36 +91,21 @@ public String getFormat() { return null; } - @Nullable - String getResponseText() { - return this.responseText; - } - /** - * @param response Chat client response to record. + * @return Chat client response * @since 1.1.0 */ - public void setResponse(ChatClientResponse response) { - if (response.chatResponse() != null) { - this.responseText = generationsToString(response.chatResponse().getResults()); - } + @Nullable + public ChatClientResponse getResponse() { + return this.response; } /** * @param response Chat client response to record. * @since 1.1.0 */ - public void appendResponse(ChatClientResponse response) { - if (this.responseText == null) { - setResponse(response); - } - else if (response.chatResponse() != null) { - this.responseText += generationsToString(response.chatResponse().getResults()); - } - } - - private String generationsToString(List generations) { - return generations.stream().map(Generation::getOutput).map(Message::getText).collect(Collectors.joining("\n")); + public void setResponse(ChatClientResponse response) { + this.response = response; } public static final class Builder { diff --git a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java index 17a803be7fe..df473e1521e 100644 --- a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java +++ b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientTests.java @@ -869,7 +869,7 @@ void whenSimplePromptThenSetRequestAndResponseOnObservationContext() { .that() .isInstanceOfSatisfying(ChatClientObservationContext.class, context -> { assertThat(context.getRequest().prompt()).isEqualTo(actualPrompt); - assertThat(context).extracting("responseText").isEqualTo("response"); + assertThat(context.getResponse()).isSameAs(chatClientResponse); }); } @@ -1415,7 +1415,8 @@ void whenSimplePromptThenSetFluxResponseOnObservationContext() { .that() .isInstanceOfSatisfying(ChatClientObservationContext.class, context -> { assertThat(context.getRequest().prompt()).isEqualTo(actualPrompt); - assertThat(context).extracting("responseText").isEqualTo("response"); + assertThat(context.getResponse().chatResponse().getResults()) + .isEqualTo(chatClientResponse.chatResponse().getResults()); }); } diff --git a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java index 0a63d392ece..59dcf6eb20d 100644 --- a/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java +++ b/spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/observation/ChatClientObservationContextTests.java @@ -190,7 +190,7 @@ void whenNoAdvisorsSpecifiedThenGetAdvisorsReturnsEmptyOrNull() { } @Test - void whenSetChatClientResponseThenReturnResponseText() { + void whenSetChatClientResponseThenReturnTheSameResponse() { var observationContext = ChatClientObservationContext.builder() .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) .build(); @@ -201,7 +201,7 @@ void whenSetChatClientResponseThenReturnResponseText() { .build(); observationContext.setResponse(response); - assertThat(observationContext.getResponseText()).isEqualTo("Test message"); + assertThat(observationContext.getResponse()).isSameAs(response); } @Test @@ -209,53 +209,9 @@ void whenSetChatClientResponseWithNullChatResponseThenReturnNull() { var observationContext = ChatClientObservationContext.builder() .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) .build(); - var nullResponse = ChatClientResponse.builder().chatResponse(null).build(); - observationContext.setResponse(nullResponse); - assertThat(observationContext.getResponseText()).isNull(); - } - - @Test - void whenAppendChatClientResponseThenReturnResponseText() { - var observationContext = ChatClientObservationContext.builder() - .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) - .build(); - var response1 = ChatClientResponse.builder() - .chatResponse(ChatResponse.builder() - .generations(List.of(new Generation(new AssistantMessage("Test message ")))) - .build()) - .build(); - var response2 = ChatClientResponse.builder() - .chatResponse(ChatResponse.builder() - .generations(List.of(new Generation(new AssistantMessage("and another test message")))) - .build()) - .build(); - - observationContext.appendResponse(response1); - assertThat(observationContext.getResponseText()).isEqualTo("Test message "); - - observationContext.appendResponse(response2); - assertThat(observationContext.getResponseText()).isEqualTo("Test message and another test message"); - } - - @Test - void whenAppendChatClientResponseWithNullChatResponseThenReturnNull() { - var observationContext = ChatClientObservationContext.builder() - .request(ChatClientRequest.builder().prompt(new Prompt("Test prompt")).build()) - .build(); - var nullResponse = ChatClientResponse.builder().chatResponse(null).build(); - var response = ChatClientResponse.builder() - .chatResponse(ChatResponse.builder() - .generations(List.of(new Generation(new AssistantMessage("Test message ")))) - .build()) - .build(); - - observationContext.appendResponse(nullResponse); - assertThat(observationContext.getResponseText()).isNull(); - observationContext.appendResponse(response); - assertThat(observationContext.getResponseText()).isEqualTo("Test message "); - observationContext.appendResponse(nullResponse); - assertThat(observationContext.getResponseText()).isEqualTo("Test message "); + observationContext.setResponse(null); + assertThat(observationContext.getResponse()).isNull(); } }