Skip to content

Commit e53bac4

Browse files
ThomasVitaleThomasVitale
committed
feat: Make advisors observations configurable
* Add possibility to provide a custom AdvisorObservationConvention to the ChatClient for customizing the conventions for advisor spans and metrics. * Add ChatClientResponse to ChatClientObservationContext for achiving full visibility into both request and response. Signed-off-by: Thomas Vitale <[email protected]>
1 parent c2103b0 commit e53bac4

File tree

7 files changed

+135
-28
lines changed

7 files changed

+135
-28
lines changed

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

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import reactor.core.publisher.Flux;
2727

2828
import org.springframework.ai.chat.client.advisor.api.Advisor;
29+
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationConvention;
2930
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
3031
import org.springframework.ai.chat.messages.Message;
3132
import org.springframework.ai.chat.model.ChatModel;
@@ -65,22 +66,46 @@ static ChatClient create(ChatModel chatModel, ObservationRegistry observationReg
6566
return create(chatModel, observationRegistry, null);
6667
}
6768

69+
/**
70+
* @deprecated in favor of
71+
* {@link #create(ChatModel, ObservationRegistry, ChatClientObservationConvention, AdvisorObservationConvention)}.
72+
*/
73+
@Deprecated
74+
static ChatClient create(ChatModel chatModel, ObservationRegistry observationRegistry,
75+
@Nullable ChatClientObservationConvention chatClientObservationConvention) {
76+
return create(chatModel, observationRegistry, chatClientObservationConvention, null);
77+
}
78+
6879
static ChatClient create(ChatModel chatModel, ObservationRegistry observationRegistry,
69-
@Nullable ChatClientObservationConvention observationConvention) {
80+
@Nullable ChatClientObservationConvention chatClientObservationConvention,
81+
@Nullable AdvisorObservationConvention advisorObservationConvention) {
7082
Assert.notNull(chatModel, "chatModel cannot be null");
7183
Assert.notNull(observationRegistry, "observationRegistry cannot be null");
72-
return builder(chatModel, observationRegistry, observationConvention).build();
84+
return builder(chatModel, observationRegistry, chatClientObservationConvention, advisorObservationConvention)
85+
.build();
7386
}
7487

7588
static Builder builder(ChatModel chatModel) {
7689
return builder(chatModel, ObservationRegistry.NOOP, null);
7790
}
7891

92+
/**
93+
* @deprecated in favor of
94+
* {@link #builder(ChatModel, ObservationRegistry, ChatClientObservationConvention, AdvisorObservationConvention)}.
95+
*/
96+
@Deprecated
97+
static Builder builder(ChatModel chatModel, ObservationRegistry observationRegistry,
98+
@Nullable ChatClientObservationConvention chatClientObservationConvention) {
99+
return builder(chatModel, observationRegistry, chatClientObservationConvention, null);
100+
}
101+
79102
static Builder builder(ChatModel chatModel, ObservationRegistry observationRegistry,
80-
@Nullable ChatClientObservationConvention customObservationConvention) {
103+
@Nullable ChatClientObservationConvention chatClientObservationConvention,
104+
@Nullable AdvisorObservationConvention advisorObservationConvention) {
81105
Assert.notNull(chatModel, "chatModel cannot be null");
82106
Assert.notNull(observationRegistry, "observationRegistry cannot be null");
83-
return new DefaultChatClientBuilder(chatModel, observationRegistry, customObservationConvention);
107+
return new DefaultChatClientBuilder(chatModel, observationRegistry, chatClientObservationConvention,
108+
advisorObservationConvention);
84109
}
85110

86111
ChatClientRequestSpec prompt();

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

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import org.springframework.ai.chat.client.advisor.DefaultAroundAdvisorChain;
3939
import org.springframework.ai.chat.client.advisor.api.Advisor;
4040
import org.springframework.ai.chat.client.advisor.api.BaseAdvisorChain;
41+
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationConvention;
4142
import org.springframework.ai.chat.client.observation.ChatClientObservationContext;
4243
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
4344
import org.springframework.ai.chat.client.observation.ChatClientObservationDocumentation;
@@ -513,7 +514,9 @@ private ChatClientResponse doGetObservableChatClientResponse(ChatClientRequest c
513514
// CHECKSTYLE:OFF
514515
var chatClientResponse = observation.observe(() -> {
515516
// Apply the advisor chain that terminates with the ChatModelCallAdvisor.
516-
return this.advisorChain.nextCall(chatClientRequest);
517+
var response = this.advisorChain.nextCall(chatClientRequest);
518+
observationContext.setChatClientResponse(response);
519+
return response;
517520
});
518521
// CHECKSTYLE:ON
519522
return chatClientResponse != null ? chatClientResponse : ChatClientResponse.builder().build();
@@ -608,7 +611,10 @@ public static class DefaultChatClientRequestSpec implements ChatClientRequestSpe
608611

609612
private final ObservationRegistry observationRegistry;
610613

611-
private final ChatClientObservationConvention observationConvention;
614+
private final ChatClientObservationConvention chatClientObservationConvention;
615+
616+
@Nullable
617+
private final AdvisorObservationConvention advisorObservationConvention;
612618

613619
private final ChatModel chatModel;
614620

@@ -649,18 +655,34 @@ public static class DefaultChatClientRequestSpec implements ChatClientRequestSpe
649655
DefaultChatClientRequestSpec(DefaultChatClientRequestSpec ccr) {
650656
this(ccr.chatModel, ccr.userText, ccr.userParams, ccr.userMetadata, ccr.systemText, ccr.systemParams,
651657
ccr.systemMetadata, ccr.toolCallbacks, ccr.messages, ccr.toolNames, ccr.media, ccr.chatOptions,
652-
ccr.advisors, ccr.advisorParams, ccr.observationRegistry, ccr.observationConvention,
653-
ccr.toolContext, ccr.templateRenderer);
658+
ccr.advisors, ccr.advisorParams, ccr.observationRegistry, ccr.chatClientObservationConvention,
659+
ccr.toolContext, ccr.templateRenderer, ccr.advisorObservationConvention);
654660
}
655661

662+
/**
663+
* @deprecated in favor of the other constructor.
664+
*/
665+
@Deprecated
656666
public DefaultChatClientRequestSpec(ChatModel chatModel, @Nullable String userText,
657667
Map<String, Object> userParams, Map<String, Object> userMetadata, @Nullable String systemText,
658668
Map<String, Object> systemParams, Map<String, Object> systemMetadata, List<ToolCallback> toolCallbacks,
659669
List<Message> messages, List<String> toolNames, List<Media> media, @Nullable ChatOptions chatOptions,
660670
List<Advisor> advisors, Map<String, Object> advisorParams, ObservationRegistry observationRegistry,
661-
@Nullable ChatClientObservationConvention observationConvention, Map<String, Object> toolContext,
662-
@Nullable TemplateRenderer templateRenderer) {
671+
@Nullable ChatClientObservationConvention chatClientObservationConvention,
672+
Map<String, Object> toolContext, @Nullable TemplateRenderer templateRenderer) {
673+
this(chatModel, userText, userParams, userMetadata, systemText, systemParams, systemMetadata, toolCallbacks,
674+
messages, toolNames, media, chatOptions, advisors, advisorParams, observationRegistry,
675+
chatClientObservationConvention, toolContext, templateRenderer, null);
676+
}
663677

678+
public DefaultChatClientRequestSpec(ChatModel chatModel, @Nullable String userText,
679+
Map<String, Object> userParams, Map<String, Object> userMetadata, @Nullable String systemText,
680+
Map<String, Object> systemParams, Map<String, Object> systemMetadata, List<ToolCallback> toolCallbacks,
681+
List<Message> messages, List<String> toolNames, List<Media> media, @Nullable ChatOptions chatOptions,
682+
List<Advisor> advisors, Map<String, Object> advisorParams, ObservationRegistry observationRegistry,
683+
@Nullable ChatClientObservationConvention chatClientObservationConvention,
684+
Map<String, Object> toolContext, @Nullable TemplateRenderer templateRenderer,
685+
@Nullable AdvisorObservationConvention advisorObservationConvention) {
664686
Assert.notNull(chatModel, "chatModel cannot be null");
665687
Assert.notNull(userParams, "userParams cannot be null");
666688
Assert.notNull(userMetadata, "userMetadata cannot be null");
@@ -694,10 +716,11 @@ public DefaultChatClientRequestSpec(ChatModel chatModel, @Nullable String userTe
694716
this.advisors.addAll(advisors);
695717
this.advisorParams.putAll(advisorParams);
696718
this.observationRegistry = observationRegistry;
697-
this.observationConvention = observationConvention != null ? observationConvention
698-
: DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION;
719+
this.chatClientObservationConvention = chatClientObservationConvention != null
720+
? chatClientObservationConvention : DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION;
699721
this.toolContext.putAll(toolContext);
700722
this.templateRenderer = templateRenderer != null ? templateRenderer : DEFAULT_TEMPLATE_RENDERER;
723+
this.advisorObservationConvention = advisorObservationConvention;
701724
}
702725

703726
@Nullable
@@ -770,7 +793,8 @@ public TemplateRenderer getTemplateRenderer() {
770793
@Override
771794
public Builder mutate() {
772795
DefaultChatClientBuilder builder = (DefaultChatClientBuilder) ChatClient
773-
.builder(this.chatModel, this.observationRegistry, this.observationConvention)
796+
.builder(this.chatModel, this.observationRegistry, this.chatClientObservationConvention,
797+
this.advisorObservationConvention)
774798
.defaultTemplateRenderer(this.templateRenderer)
775799
.defaultToolCallbacks(this.toolCallbacks)
776800
.defaultToolContext(this.toolContext)
@@ -990,14 +1014,14 @@ public ChatClientRequestSpec templateRenderer(TemplateRenderer templateRenderer)
9901014
public CallResponseSpec call() {
9911015
BaseAdvisorChain advisorChain = buildAdvisorChain();
9921016
return new DefaultCallResponseSpec(DefaultChatClientUtils.toChatClientRequest(this), advisorChain,
993-
this.observationRegistry, this.observationConvention);
1017+
this.observationRegistry, this.chatClientObservationConvention);
9941018
}
9951019

9961020
@Override
9971021
public StreamResponseSpec stream() {
9981022
BaseAdvisorChain advisorChain = buildAdvisorChain();
9991023
return new DefaultStreamResponseSpec(DefaultChatClientUtils.toChatClientRequest(this), advisorChain,
1000-
this.observationRegistry, this.observationConvention);
1024+
this.observationRegistry, this.chatClientObservationConvention);
10011025
}
10021026

10031027
private BaseAdvisorChain buildAdvisorChain() {
@@ -1006,7 +1030,10 @@ private BaseAdvisorChain buildAdvisorChain() {
10061030
this.advisors.add(ChatModelCallAdvisor.builder().chatModel(this.chatModel).build());
10071031
this.advisors.add(ChatModelStreamAdvisor.builder().chatModel(this.chatModel).build());
10081032

1009-
return DefaultAroundAdvisorChain.builder(this.observationRegistry).pushAll(this.advisors).build();
1033+
return DefaultAroundAdvisorChain.builder(this.observationRegistry)
1034+
.observationConvention(this.advisorObservationConvention)
1035+
.pushAll(this.advisors)
1036+
.build();
10101037
}
10111038

10121039
}

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.springframework.ai.chat.client.ChatClient.PromptUserSpec;
3030
import org.springframework.ai.chat.client.DefaultChatClient.DefaultChatClientRequestSpec;
3131
import org.springframework.ai.chat.client.advisor.api.Advisor;
32+
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationConvention;
3233
import org.springframework.ai.chat.client.observation.ChatClientObservationConvention;
3334
import org.springframework.ai.chat.messages.Message;
3435
import org.springframework.ai.chat.model.ChatModel;
@@ -60,13 +61,24 @@ public class DefaultChatClientBuilder implements Builder {
6061
this(chatModel, ObservationRegistry.NOOP, null);
6162
}
6263

64+
/**
65+
* @deprecated in favor of
66+
* {@link #DefaultChatClientBuilder(ChatModel, ObservationRegistry, ChatClientObservationConvention, AdvisorObservationConvention)}.
67+
*/
68+
@Deprecated
6369
public DefaultChatClientBuilder(ChatModel chatModel, ObservationRegistry observationRegistry,
64-
@Nullable ChatClientObservationConvention customObservationConvention) {
70+
@Nullable ChatClientObservationConvention chatClientObservationConvention) {
71+
this(chatModel, observationRegistry, chatClientObservationConvention, null);
72+
}
73+
74+
public DefaultChatClientBuilder(ChatModel chatModel, ObservationRegistry observationRegistry,
75+
@Nullable ChatClientObservationConvention chatClientObservationConvention,
76+
@Nullable AdvisorObservationConvention advisorObservationConvention) {
6577
Assert.notNull(chatModel, "the " + ChatModel.class.getName() + " must be non-null");
6678
Assert.notNull(observationRegistry, "the " + ObservationRegistry.class.getName() + " must be non-null");
6779
this.defaultRequest = new DefaultChatClientRequestSpec(chatModel, null, Map.of(), Map.of(), null, Map.of(),
6880
Map.of(), List.of(), List.of(), List.of(), List.of(), null, List.of(), Map.of(), observationRegistry,
69-
customObservationConvention, Map.of(), null);
81+
chatClientObservationConvention, Map.of(), null, advisorObservationConvention);
7082
}
7183

7284
public ChatClient build() {

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

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,8 @@
3535
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationConvention;
3636
import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationDocumentation;
3737
import org.springframework.ai.chat.client.advisor.observation.DefaultAdvisorObservationConvention;
38-
import org.springframework.ai.template.TemplateRenderer;
39-
import org.springframework.ai.template.st.StTemplateRenderer;
4038
import org.springframework.core.OrderComparator;
39+
import org.springframework.lang.Nullable;
4140
import org.springframework.util.Assert;
4241
import org.springframework.util.CollectionUtils;
4342

@@ -55,8 +54,6 @@ public class DefaultAroundAdvisorChain implements BaseAdvisorChain {
5554

5655
public static final AdvisorObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultAdvisorObservationConvention();
5756

58-
private static final TemplateRenderer DEFAULT_TEMPLATE_RENDERER = StTemplateRenderer.builder().build();
59-
6057
private final List<CallAdvisor> originalCallAdvisors;
6158

6259
private final List<StreamAdvisor> originalStreamAdvisors;
@@ -67,8 +64,10 @@ public class DefaultAroundAdvisorChain implements BaseAdvisorChain {
6764

6865
private final ObservationRegistry observationRegistry;
6966

67+
private final AdvisorObservationConvention observationConvention;
68+
7069
DefaultAroundAdvisorChain(ObservationRegistry observationRegistry, Deque<CallAdvisor> callAdvisors,
71-
Deque<StreamAdvisor> streamAdvisors) {
70+
Deque<StreamAdvisor> streamAdvisors, @Nullable AdvisorObservationConvention observationConvention) {
7271

7372
Assert.notNull(observationRegistry, "the observationRegistry must be non-null");
7473
Assert.notNull(callAdvisors, "the callAdvisors must be non-null");
@@ -79,6 +78,8 @@ public class DefaultAroundAdvisorChain implements BaseAdvisorChain {
7978
this.streamAdvisors = streamAdvisors;
8079
this.originalCallAdvisors = List.copyOf(callAdvisors);
8180
this.originalStreamAdvisors = List.copyOf(streamAdvisors);
81+
this.observationConvention = observationConvention != null ? observationConvention
82+
: DEFAULT_OBSERVATION_CONVENTION;
8283
}
8384

8485
public static Builder builder(ObservationRegistry observationRegistry) {
@@ -102,8 +103,13 @@ public ChatClientResponse nextCall(ChatClientRequest chatClientRequest) {
102103
.build();
103104

104105
return AdvisorObservationDocumentation.AI_ADVISOR
105-
.observation(null, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, this.observationRegistry)
106-
.observe(() -> advisor.adviseCall(chatClientRequest, this));
106+
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
107+
this.observationRegistry)
108+
.observe(() -> {
109+
var chatClientResponse = advisor.adviseCall(chatClientRequest, this);
110+
observationContext.setChatClientResponse(chatClientResponse);
111+
return chatClientResponse;
112+
});
107113
}
108114

109115
@Override
@@ -123,7 +129,7 @@ public Flux<ChatClientResponse> nextStream(ChatClientRequest chatClientRequest)
123129
.order(advisor.getOrder())
124130
.build();
125131

126-
var observation = AdvisorObservationDocumentation.AI_ADVISOR.observation(null,
132+
var observation = AdvisorObservationDocumentation.AI_ADVISOR.observation(this.observationConvention,
127133
DEFAULT_OBSERVATION_CONVENTION, () -> observationContext, this.observationRegistry);
128134

129135
observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
@@ -160,12 +166,20 @@ public static class Builder {
160166

161167
private final Deque<StreamAdvisor> streamAdvisors;
162168

169+
@Nullable
170+
private AdvisorObservationConvention observationConvention;
171+
163172
public Builder(ObservationRegistry observationRegistry) {
164173
this.observationRegistry = observationRegistry;
165174
this.callAdvisors = new ConcurrentLinkedDeque<>();
166175
this.streamAdvisors = new ConcurrentLinkedDeque<>();
167176
}
168177

178+
public Builder observationConvention(@Nullable AdvisorObservationConvention observationConvention) {
179+
this.observationConvention = observationConvention;
180+
return this;
181+
}
182+
169183
public Builder push(Advisor advisor) {
170184
Assert.notNull(advisor, "the advisor must be non-null");
171185
return this.pushAll(List.of(advisor));
@@ -214,7 +228,8 @@ private void reOrder() {
214228
}
215229

216230
public DefaultAroundAdvisorChain build() {
217-
return new DefaultAroundAdvisorChain(this.observationRegistry, this.callAdvisors, this.streamAdvisors);
231+
return new DefaultAroundAdvisorChain(this.observationRegistry, this.callAdvisors, this.streamAdvisors,
232+
this.observationConvention);
218233
}
219234

220235
}

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

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
import org.springframework.ai.chat.client.ChatClientAttributes;
2424
import org.springframework.ai.chat.client.ChatClientRequest;
25+
import org.springframework.ai.chat.client.ChatClientResponse;
2526
import org.springframework.ai.chat.client.advisor.api.Advisor;
2627
import org.springframework.ai.observation.AiOperationMetadata;
2728
import org.springframework.ai.observation.conventions.AiOperationType;
@@ -48,6 +49,9 @@ public class ChatClientObservationContext extends Observation.Context {
4849

4950
private final boolean stream;
5051

52+
@Nullable
53+
private ChatClientResponse chatClientResponse;
54+
5155
ChatClientObservationContext(ChatClientRequest chatClientRequest, List<? extends Advisor> advisors,
5256
boolean isStream) {
5357
Assert.notNull(chatClientRequest, "chatClientRequest cannot be null");
@@ -78,6 +82,15 @@ public boolean isStream() {
7882
return this.stream;
7983
}
8084

85+
@Nullable
86+
public ChatClientResponse getChatClientResponse() {
87+
return this.chatClientResponse;
88+
}
89+
90+
public void setChatClientResponse(@Nullable ChatClientResponse chatClientResponse) {
91+
this.chatClientResponse = chatClientResponse;
92+
}
93+
8194
@Nullable
8295
public String getFormat() {
8396
if (this.request.context().get(ChatClientAttributes.OUTPUT_FORMAT.getKey()) instanceof String format) {

‎spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/DefaultChatClientBuilderTests.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2024 the original author or authors.
2+
* Copyright 2023-2025 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
1818

1919
import java.nio.charset.Charset;
2020

21+
import io.micrometer.observation.ObservationRegistry;
2122
import org.junit.jupiter.api.Test;
2223

2324
import org.springframework.ai.chat.model.ChatModel;
@@ -63,6 +64,12 @@ void whenObservationRegistryIsNullThenThrows() {
6364
.hasMessage("the io.micrometer.observation.ObservationRegistry must be non-null");
6465
}
6566

67+
@Test
68+
void whenAdvisorObservationConventionIsNullThenReturn() {
69+
var builder = new DefaultChatClientBuilder(mock(ChatModel.class), mock(ObservationRegistry.class), null, null);
70+
assertThat(builder).isNotNull();
71+
}
72+
6673
@Test
6774
void whenUserResourceIsNullThenThrows() {
6875
DefaultChatClientBuilder builder = new DefaultChatClientBuilder(mock(ChatModel.class));

‎spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/advisor/DefaultAroundAdvisorChainTests.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,14 @@ void whenAdvisorListContainsNullElementsThenThrow() {
7676
.hasMessage("the advisors must not contain null elements");
7777
}
7878

79+
@Test
80+
void getObservationConventionIsNullThenUseDefault() {
81+
AdvisorChain chain = DefaultAroundAdvisorChain.builder(ObservationRegistry.create())
82+
.observationConvention(null)
83+
.build();
84+
assertThat(chain).isNotNull();
85+
}
86+
7987
@Test
8088
void getObservationRegistry() {
8189
ObservationRegistry observationRegistry = ObservationRegistry.create();

0 commit comments

Comments
 (0)