Skip to content

Commit 13de219

Browse files
jonatan-ivanovsobychacko
jonatan-ivanov
authored and
sobychacko
committed
Log chat client completion data and make the response text available in the Observation Context
 - Add ChatClientCompletionObservationHandler to log completion data - Utilize ChatClientMessageAggregator to set the response in the observation context - Configure observation handling via ChatClientAutoConfiguration - Testing changes and docs
1 parent 357190d commit 13de219

File tree

11 files changed

+486
-14
lines changed

11 files changed

+486
-14
lines changed

‎auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientAutoConfiguration.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
import org.springframework.ai.chat.client.ChatClient;
2525
import org.springframework.ai.chat.client.ChatClientCustomizer;
26+
import org.springframework.ai.chat.client.observation.ChatClientCompletionObservationHandler;
2627
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
2728
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
2829
import org.springframework.ai.chat.client.observation.ChatClientPromptContentObservationHandler;
@@ -71,6 +72,11 @@ private static void logPromptContentWarning() {
7172
"You have enabled logging out the ChatClient prompt content with the risk of exposing sensitive or private information. Please, be careful!");
7273
}
7374

75+
private static void logCompletionWarning() {
76+
logger.warn(
77+
"You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!");
78+
}
79+
7480
@Bean
7581
@ConditionalOnMissingBean
7682
ChatClientBuilderConfigurer chatClientBuilderConfigurer(ObjectProvider<ChatClientCustomizer> customizerProvider) {
@@ -108,6 +114,17 @@ TracingAwareLoggingObservationHandler<ChatClientObservationContext> chatClientPr
108114
return new TracingAwareLoggingObservationHandler<>(new ChatClientPromptContentObservationHandler(), tracer);
109115
}
110116

117+
@Bean
118+
@ConditionalOnMissingBean(value = ChatClientCompletionObservationHandler.class,
119+
name = "chatClientCompletionObservationHandler")
120+
@ConditionalOnProperty(prefix = ChatClientBuilderProperties.CONFIG_PREFIX + ".observations",
121+
name = "log-completion", havingValue = "true")
122+
TracingAwareLoggingObservationHandler<ChatClientObservationContext> chatClientCompletionObservationHandler(
123+
Tracer tracer) {
124+
logCompletionWarning();
125+
return new TracingAwareLoggingObservationHandler<>(new ChatClientCompletionObservationHandler(), tracer);
126+
}
127+
111128
}
112129

113130
@Configuration(proxyBeanMethods = false)
@@ -123,6 +140,15 @@ ChatClientPromptContentObservationHandler chatClientPromptContentObservationHand
123140
return new ChatClientPromptContentObservationHandler();
124141
}
125142

143+
@Bean
144+
@ConditionalOnMissingBean
145+
@ConditionalOnProperty(prefix = ChatClientBuilderProperties.CONFIG_PREFIX + ".observations",
146+
name = "log-completion", havingValue = "true")
147+
ChatClientCompletionObservationHandler chatClientCompletionObservationHandler() {
148+
logCompletionWarning();
149+
return new ChatClientCompletionObservationHandler();
150+
}
151+
126152
}
127153

128154
}

‎auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/main/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientBuilderProperties.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
* @author Josh Long
2727
* @author Arjen Poutsma
2828
* @author Thomas Vitale
29+
* @author Jonatan Ivanov
2930
* @since 1.0.0
3031
*/
3132
@ConfigurationProperties(ChatClientBuilderProperties.CONFIG_PREFIX)
@@ -59,14 +60,36 @@ public static class Observations {
5960
*/
6061
private boolean logPrompt = false;
6162

63+
/**
64+
* Whether to log the completion content in the observations.
65+
* @since 1.1.0
66+
*/
67+
private boolean logCompletion = false;
68+
6269
public boolean isLogPrompt() {
6370
return this.logPrompt;
6471
}
6572

73+
/**
74+
* @return Whether logging completion data is enabled or not.
75+
* @since 1.1.0
76+
*/
77+
public boolean isLogCompletion() {
78+
return this.logCompletion;
79+
}
80+
6681
public void setLogPrompt(boolean logPrompt) {
6782
this.logPrompt = logPrompt;
6883
}
6984

85+
/**
86+
* @param logCompletion should completion data logging be enabled or not.
87+
* @since 1.1.0
88+
*/
89+
public void setLogCompletion(boolean logCompletion) {
90+
this.logCompletion = logCompletion;
91+
}
92+
7093
}
7194

7295
}

‎auto-configurations/models/chat/client/spring-ai-autoconfigure-model-chat-client/src/test/java/org/springframework/ai/model/chat/client/autoconfigure/ChatClientObservationAutoConfigurationTests.java‎

Lines changed: 120 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.junit.jupiter.api.Test;
2121
import org.junit.jupiter.api.extension.ExtendWith;
2222

23+
import org.springframework.ai.chat.client.observation.ChatClientCompletionObservationHandler;
2324
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
2425
import org.springframework.ai.chat.client.observation.ChatClientPromptContentObservationHandler;
2526
import org.springframework.ai.observation.TracingAwareLoggingObservationHandler;
@@ -48,16 +49,18 @@ class ChatClientObservationAutoConfigurationTests {
4849
.withConfiguration(AutoConfigurations.of(ChatClientAutoConfiguration.class));
4950

5051
@Test
51-
void promptContentHandlerNoTracer() {
52+
void handlersNoTracer() {
5253
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
5354
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
55+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
5456
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
5557
}
5658

5759
@Test
58-
void promptContentHandlerWithTracer() {
60+
void handlersWithTracer() {
5961
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
6062
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
63+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
6164
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
6265
}
6366

@@ -66,6 +69,7 @@ void promptContentHandlerEnabledNoTracer(CapturedOutput output) {
6669
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
6770
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")
6871
.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)
72+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
6973
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
7074
assertThat(output).contains(
7175
"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) {
7680
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
7781
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")
7882
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
83+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
7984
.hasSingleBean(TracingAwareLoggingObservationHandler.class));
8085
assertThat(output).contains(
8186
"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() {
8691
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
8792
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=false")
8893
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
94+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
8995
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
9096
}
9197

@@ -94,6 +100,47 @@ void promptContentHandlerDisabledWithTracer() {
94100
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
95101
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=false")
96102
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
103+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
104+
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
105+
}
106+
107+
@Test
108+
void completionHandlerEnabledNoTracer(CapturedOutput output) {
109+
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
110+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=true")
111+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
112+
.hasSingleBean(ChatClientCompletionObservationHandler.class)
113+
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
114+
assertThat(output).contains(
115+
"You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!");
116+
}
117+
118+
@Test
119+
void completionHandlerEnabledWithTracer(CapturedOutput output) {
120+
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
121+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=true")
122+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
123+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
124+
.hasSingleBean(TracingAwareLoggingObservationHandler.class));
125+
assertThat(output).contains(
126+
"You have enabled logging out the ChatClient completion content with the risk of exposing sensitive or private information. Please, be careful!");
127+
}
128+
129+
@Test
130+
void completionHandlerDisabledNoTracer() {
131+
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
132+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=false")
133+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
134+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
135+
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
136+
}
137+
138+
@Test
139+
void completionDisabledWithTracer() {
140+
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
141+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=false")
142+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
143+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
97144
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
98145
}
99146

@@ -104,6 +151,7 @@ void customChatClientPromptContentObservationHandlerNoTracer() {
104151
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")
105152
.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)
106153
.hasBean("customChatClientPromptContentObservationHandler")
154+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
107155
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
108156
}
109157

@@ -114,20 +162,61 @@ void customChatClientPromptContentObservationHandlerWithTracer() {
114162
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")
115163
.run(context -> assertThat(context).hasSingleBean(ChatClientPromptContentObservationHandler.class)
116164
.hasBean("customChatClientPromptContentObservationHandler")
165+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class)
117166
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
118167
}
119168

120169
@Test
121-
void customTracingAwareLoggingObservationHandler() {
170+
void customTracingAwareLoggingObservationHandlerForChatClientPromptContent() {
122171
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
123-
.withUserConfiguration(CustomTracingAwareLoggingObservationHandlerConfiguration.class)
172+
.withUserConfiguration(
173+
CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration.class)
124174
.withPropertyValues("spring.ai.chat.client.observations.log-prompt=true")
125175
.run(context -> {
126176
assertThat(context).hasSingleBean(TracingAwareLoggingObservationHandler.class)
127177
.hasBean("chatClientPromptContentObservationHandler")
128-
.doesNotHaveBean(ChatClientPromptContentObservationHandler.class);
129-
assertThat(context.getBean(TracingAwareLoggingObservationHandler.class))
130-
.isSameAs(CustomTracingAwareLoggingObservationHandlerConfiguration.handlerInstance);
178+
.doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
179+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class);
180+
assertThat(context.getBean(TracingAwareLoggingObservationHandler.class)).isSameAs(
181+
CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration.handlerInstance);
182+
});
183+
}
184+
185+
@Test
186+
void customChatClientCompletionObservationHandlerNoTracer() {
187+
this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))
188+
.withUserConfiguration(CustomChatClientCompletionObservationHandlerConfiguration.class)
189+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=true")
190+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
191+
.hasSingleBean(ChatClientCompletionObservationHandler.class)
192+
.hasBean("customChatClientCompletionObservationHandler")
193+
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
194+
}
195+
196+
@Test
197+
void customChatClientCompletionObservationHandlerWithTracer() {
198+
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
199+
.withUserConfiguration(CustomChatClientCompletionObservationHandlerConfiguration.class)
200+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=true")
201+
.run(context -> assertThat(context).doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
202+
.hasSingleBean(ChatClientCompletionObservationHandler.class)
203+
.hasBean("customChatClientCompletionObservationHandler")
204+
.doesNotHaveBean(TracingAwareLoggingObservationHandler.class));
205+
}
206+
207+
@Test
208+
void customTracingAwareLoggingObservationHandlerForChatClientCompletion() {
209+
this.contextRunner.withUserConfiguration(TracerConfiguration.class)
210+
.withUserConfiguration(
211+
CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration.class)
212+
.withPropertyValues("spring.ai.chat.client.observations.log-completion=true")
213+
.run(context -> {
214+
assertThat(context).hasSingleBean(TracingAwareLoggingObservationHandler.class)
215+
.hasBean("chatClientCompletionObservationHandler")
216+
.doesNotHaveBean(ChatClientPromptContentObservationHandler.class)
217+
.doesNotHaveBean(ChatClientCompletionObservationHandler.class);
218+
assertThat(context.getBean(TracingAwareLoggingObservationHandler.class)).isSameAs(
219+
CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration.handlerInstance);
131220
});
132221
}
133222

@@ -152,7 +241,7 @@ ChatClientPromptContentObservationHandler customChatClientPromptContentObservati
152241
}
153242

154243
@Configuration(proxyBeanMethods = false)
155-
static class CustomTracingAwareLoggingObservationHandlerConfiguration {
244+
static class CustomTracingAwareLoggingObservationHandlerForChatClientPromptContentConfiguration {
156245

157246
static TracingAwareLoggingObservationHandler<ChatClientObservationContext> handlerInstance = new TracingAwareLoggingObservationHandler<>(
158247
new ChatClientPromptContentObservationHandler(), null);
@@ -164,4 +253,27 @@ TracingAwareLoggingObservationHandler<ChatClientObservationContext> chatClientPr
164253

165254
}
166255

256+
@Configuration(proxyBeanMethods = false)
257+
static class CustomChatClientCompletionObservationHandlerConfiguration {
258+
259+
@Bean
260+
ChatClientCompletionObservationHandler customChatClientCompletionObservationHandler() {
261+
return new ChatClientCompletionObservationHandler();
262+
}
263+
264+
}
265+
266+
@Configuration(proxyBeanMethods = false)
267+
static class CustomTracingAwareLoggingObservationHandlerForChatClientChatClientCompletionConfiguration {
268+
269+
static TracingAwareLoggingObservationHandler<ChatClientObservationContext> handlerInstance = new TracingAwareLoggingObservationHandler<>(
270+
new ChatClientCompletionObservationHandler(), null);
271+
272+
@Bean
273+
TracingAwareLoggingObservationHandler<ChatClientObservationContext> chatClientCompletionObservationHandler() {
274+
return handlerInstance;
275+
}
276+
277+
}
278+
167279
}

‎spring-ai-client-chat/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,12 @@
110110
<scope>test</scope>
111111
</dependency>
112112

113+
<dependency>
114+
<groupId>io.micrometer</groupId>
115+
<artifactId>micrometer-observation-test</artifactId>
116+
<scope>test</scope>
117+
</dependency>
118+
113119
<dependency>
114120
<groupId>com.fasterxml.jackson.module</groupId>
115121
<artifactId>jackson-module-kotlin</artifactId>

‎spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
* @author Soby Chacko
7777
* @author Dariusz Jedrzejczyk
7878
* @author Thomas Vitale
79+
* @author Jonatan Ivanov
7980
* @since 1.0.0
8081
*/
8182
public class DefaultChatClient implements ChatClient {
@@ -84,6 +85,8 @@ public class DefaultChatClient implements ChatClient {
8485

8586
private static final TemplateRenderer DEFAULT_TEMPLATE_RENDERER = StTemplateRenderer.builder().build();
8687

88+
private static final ChatClientMessageAggregator CHAT_CLIENT_MESSAGE_AGGREGATOR = new ChatClientMessageAggregator();
89+
8790
private final DefaultChatClientRequestSpec defaultChatClientRequest;
8891

8992
public DefaultChatClient(DefaultChatClientRequestSpec defaultChatClientRequest) {
@@ -513,7 +516,9 @@ private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest c
513516
// CHECKSTYLE:OFF
514517
var chatClientResponse = observation.observe(() -> {
515518
// Apply the advisor chain that terminates with the ChatModelCallAdvisor.
516-
return this.advisorChain.nextCall(chatClientRequest);
519+
var response = this.advisorChain.nextCall(chatClientRequest);
520+
observationContext.setResponse(response);
521+
return response;
517522
});
518523
// CHECKSTYLE:ON
519524
return chatClientResponse != null ? chatClientResponse : ChatClientResponse.builder().build();
@@ -571,11 +576,13 @@ private Flux<ChatClientResponse> doGetObservableFluxChatResponse(ChatClientReque
571576

572577
// @formatter:off
573578
// Apply the advisor chain that terminates with the ChatModelStreamAdvisor.
574-
return this.advisorChain.nextStream(chatClientRequest)
579+
Flux<ChatClientResponse> chatClientResponse = this.advisorChain.nextStream(chatClientRequest)
575580
.doOnError(observation::error)
576581
.doFinally(s -> observation.stop())
577582
.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
578583
// @formatter:on
584+
return CHAT_CLIENT_MESSAGE_AGGREGATOR.aggregateChatClientResponse(chatClientResponse,
585+
observationContext::setResponse);
579586
});
580587
}
581588

0 commit comments

Comments
 (0)