Improve chat client and advisor observations

* Make ChatClient and Advisor observation logic null-safe
* Simplify naming for Advisor observations
* Include high-cardinality attributes only if a value is present
* Fix condition to include system test to chat client observations
* Add Advisor order information to context
* Streamline usage of enums and utils to reduce hard-coded/duplications
* Fix pending Chroma integration test

Signed-off-by: Thomas Vitale <[email protected]>

Author: ThomasVitale

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

@@ -77,6 +77,7 @@
  * @author Arjen Poutsma
  * @author Soby Chacko
  * @author Dariusz Jedrzejczyk
+ * @author Thomas Vitale
  * @since 1.0.0
  */
 public class DefaultChatClient implements ChatClient {
@@ -360,8 +361,11 @@ private ChatResponse doGetChatResponse() {
 		private ChatResponse doGetObservableChatResponse(DefaultChatClientRequestSpec inputRequest,
 				String formatParam) {
 
-			ChatClientObservationContext observationContext = new ChatClientObservationContext(inputRequest,
-					formatParam, false);
+			ChatClientObservationContext observationContext = ChatClientObservationContext.builder()
+				.withRequest(inputRequest)
+				.withFormat(formatParam)
+				.withStream(false)
+				.build();
 
 			var observation = ChatClientObservationDocumentation.AI_CHAT_CLIENT.observation(
 					inputRequest.getCustomObservationConvention(), DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION,
@@ -407,8 +411,10 @@ public DefaultStreamResponseSpec(DefaultChatClientRequestSpec request) {
 		private Flux<ChatResponse> doGetObservableFluxChatResponse(DefaultChatClientRequestSpec inputRequest) {
 			return Flux.deferContextual(contextView -> {
 
-				ChatClientObservationContext observationContext = new ChatClientObservationContext(inputRequest, "",
-						true);
+				ChatClientObservationContext observationContext = ChatClientObservationContext.builder()
+					.withRequest(inputRequest)
+					.withStream(true)
+					.build();
 
 				Observation observation = ChatClientObservationDocumentation.AI_CHAT_CLIENT.observation(
 						inputRequest.getCustomObservationConvention(), DEFAULT_CHAT_CLIENT_OBSERVATION_CONVENTION,

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

@@ -85,6 +85,7 @@ public AdvisedResponse nextAroundCall(AdvisedRequest advisedRequest) {
 			.withAdvisorType(AdvisorObservationContext.Type.AROUND)
 			.withAdvisedRequest(advisedRequest)
 			.withAdvisorRequestContext(advisedRequest.adviseContext())
+			.withOrder(advisor.getOrder())
 			.build();
 
 		return AdvisorObservationDocumentation.AI_ADVISOR
@@ -106,6 +107,7 @@ public Flux<AdvisedResponse> nextAroundStream(AdvisedRequest advisedRequest) {
 				.withAdvisorType(AdvisorObservationContext.Type.AROUND)
 				.withAdvisedRequest(advisedRequest)
 				.withAdvisorRequestContext(advisedRequest.adviseContext())
+				.withOrder(advisor.getOrder())
 				.build();
 
 			var observation = AdvisorObservationDocumentation.AI_ADVISOR.observation(null,

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/observation/AdvisorObservationContext.java

@@ -20,15 +20,18 @@
 import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
 import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 import io.micrometer.observation.Observation;
 
 /**
+ * Context used to store metadata for chat client advisors.
+ *
  * @author Christian Tzolov
+ * @author Thomas Vitale
  * @since 1.0.0
  */
-
 public class AdvisorObservationContext extends Observation.Context {
 
 	public enum Type {
@@ -37,35 +40,48 @@ public enum Type {
 
 	}
 
-	private String advisorName;
+	private final String advisorName;
 
-	private Type advisorType;
+	private final Type advisorType;
 
 	/**
 	 * The {@link AdvisedRequest} data to be advised. Represents the row
 	 * {@link ChatClient.ChatClientRequestSpec} data before sealed into a {@link Prompt}.
 	 */
+	@Nullable
 	private AdvisedRequest advisorRequest;
 
 	/**
 	 * The shared data between the advisors in the chain. It is shared between all request
 	 * and response advising points of all advisors in the chain.
 	 */
+	@Nullable
 	private Map<String, Object> advisorRequestContext;
 
 	/**
 	 * the shared data between the advisors in the chain. It is shared between all request
 	 * and response advising points of all advisors in the chain.
 	 */
+	@Nullable
 	private Map<String, Object> advisorResponseContext;
 
 	/**
 	 * The order of the advisor in the advisor chain.
 	 */
-	private int order;
+	private final int order;
+
+	public AdvisorObservationContext(String advisorName, Type advisorType, @Nullable AdvisedRequest advisorRequest,
+			@Nullable Map<String, Object> advisorRequestContext, @Nullable Map<String, Object> advisorResponseContext,
+			int order) {
+		Assert.hasText(advisorName, "advisorName must not be null or empty");
+		Assert.notNull(advisorType, "advisorType must not be null");
 
-	public void setAdvisorName(String advisorName) {
 		this.advisorName = advisorName;
+		this.advisorType = advisorType;
+		this.advisorRequest = advisorRequest;
+		this.advisorRequestContext = advisorRequestContext;
+		this.advisorResponseContext = advisorResponseContext;
+		this.order = order;
 	}
 
 	public String getAdvisorName() {
@@ -76,84 +92,88 @@ public Type getAdvisorType() {
 		return this.advisorType;
 	}
 
-	public void setAdvisorType(Type type) {
-		this.advisorType = type;
-	}
-
+	@Nullable
 	public AdvisedRequest getAdvisedRequest() {
 		return this.advisorRequest;
 	}
 
-	public void setAdvisedRequest(AdvisedRequest advisedRequest) {
+	public void setAdvisedRequest(@Nullable AdvisedRequest advisedRequest) {
 		this.advisorRequest = advisedRequest;
 	}
 
+	@Nullable
 	public Map<String, Object> getAdvisorRequestContext() {
 		return this.advisorRequestContext;
 	}
 
-	public void setAdvisorRequestContext(Map<String, Object> advisorRequestContext) {
+	public void setAdvisorRequestContext(@Nullable Map<String, Object> advisorRequestContext) {
 		this.advisorRequestContext = advisorRequestContext;
 	}
 
+	@Nullable
 	public Map<String, Object> getAdvisorResponseContext() {
 		return this.advisorResponseContext;
 	}
 
-	public void setAdvisorResponseContext(Map<String, Object> advisorResponseContext) {
+	public void setAdvisorResponseContext(@Nullable Map<String, Object> advisorResponseContext) {
 		this.advisorResponseContext = advisorResponseContext;
 	}
 
 	public int getOrder() {
 		return this.order;
 	}
 
-	public void setOrder(int order) {
-		this.order = order;
-	}
-
 	public static Builder builder() {
 		return new Builder();
 	}
 
 	public static class Builder {
 
-		private final AdvisorObservationContext context = new AdvisorObservationContext();
+		private String advisorName;
+
+		private Type advisorType;
+
+		private AdvisedRequest advisorRequest;
+
+		private Map<String, Object> advisorRequestContext;
+
+		private Map<String, Object> advisorResponseContext;
+
+		private int order = 0;
 
 		public Builder withAdvisorName(String advisorName) {
-			this.context.setAdvisorName(advisorName);
+			this.advisorName = advisorName;
 			return this;
 		}
 
 		public Builder withAdvisorType(Type advisorType) {
-			this.context.setAdvisorType(advisorType);
+			this.advisorType = advisorType;
 			return this;
 		}
 
 		public Builder withAdvisedRequest(AdvisedRequest advisedRequest) {
-			this.context.setAdvisedRequest(advisedRequest);
+			this.advisorRequest = advisedRequest;
 			return this;
 		}
 
 		public Builder withAdvisorRequestContext(Map<String, Object> advisorRequestContext) {
-			this.context.setAdvisorRequestContext(advisorRequestContext);
+			this.advisorRequestContext = advisorRequestContext;
 			return this;
 		}
 
 		public Builder withAdvisorResponseContext(Map<String, Object> advisorResponseContext) {
-			this.context.setAdvisorResponseContext(advisorResponseContext);
+			this.advisorResponseContext = advisorResponseContext;
 			return this;
 		}
 
 		public Builder withOrder(int order) {
-			this.context.setOrder(order);
+			this.order = order;
 			return this;
 		}
 
 		public AdvisorObservationContext build() {
-			Assert.hasText(this.context.advisorName, "The advisorName must not be empty!");
-			Assert.notNull(this.context.advisorType, "The advisorType must not be null!");
-			return this.context;
+			return new AdvisorObservationContext(advisorName, advisorType, advisorRequest, advisorRequestContext,
+					advisorResponseContext, order);
 		}
 
 	}

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/observation/AdvisorObservationConvention.java

@@ -19,6 +19,8 @@
 import io.micrometer.observation.ObservationConvention;
 
 /**
+ * Interface for an {@link ObservationConvention} for chat client advisors.
+ *
  * @author Christian Tzolov
  * @since 1.0.0
  */

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/observation/AdvisorObservationDocumentation.java

@@ -27,7 +27,7 @@
 public enum AdvisorObservationDocumentation implements ObservationDocumentation {
 
 	/**
-	 * AI Chat Client observations
+	 * AI Advisor observations
 	 */
 	AI_ADVISOR {
 		@Override
@@ -65,7 +65,7 @@ public String asString() {
 		ADVISOR_TYPE {
 			@Override
 			public String asString() {
-				return "spring.ai.chat.client.advisor.type";
+				return "spring.ai.advisor.type";
 			}
 		}
 
@@ -74,12 +74,12 @@ public String asString() {
 	public enum HighCardinalityKeyNames implements KeyName {
 
 		/**
-		 * Chat Model name.
+		 * Advisor name.
 		 */
 		ADVISOR_NAME {
 			@Override
 			public String asString() {
-				return "spring.ai.chat.client.advisor.name";
+				return "spring.ai.advisor.name";
 			}
 		},
 		/**
@@ -88,7 +88,7 @@ public String asString() {
 		ADVISOR_ORDER {
 			@Override
 			public String asString() {
-				return "spring.ai.chat.client.advisor.order";
+				return "spring.ai.advisor.order";
 			}
 		}
 

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/observation/DefaultAdvisorObservationConvention.java

@@ -17,6 +17,7 @@
 
 import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationDocumentation.HighCardinalityKeyNames;
 import org.springframework.ai.chat.client.advisor.observation.AdvisorObservationDocumentation.LowCardinalityKeyNames;
+import org.springframework.ai.observation.conventions.SpringAiKind;
 import org.springframework.ai.util.ParsingUtils;
 import org.springframework.lang.Nullable;
 
@@ -27,18 +28,9 @@
  * @author Christian Tzolov
  * @since 1.0.0
  */
-
 public class DefaultAdvisorObservationConvention implements AdvisorObservationConvention {
 
-	public static final String DEFAULT_NAME = "spring.ai.chat.client.advisor";
-
-	private static final String CHAT_CLIENT_ADVISOR_SPRING_AI_KIND = "chat_client_advisor";
-
-	private static final KeyValue ADVISOR_TYPE_NONE = KeyValue.of(LowCardinalityKeyNames.ADVISOR_TYPE,
-			KeyValue.NONE_VALUE);
-
-	private static final KeyValue ADVISOR_NAME_NONE = KeyValue.of(HighCardinalityKeyNames.ADVISOR_NAME,
-			KeyValue.NONE_VALUE);
+	public static final String DEFAULT_NAME = "spring.ai.advisor";
 
 	private final String name;
 
@@ -73,15 +65,12 @@ public KeyValues getLowCardinalityKeyValues(AdvisorObservationContext context) {
 	}
 
 	protected KeyValue advisorType(AdvisorObservationContext context) {
-		if (context.getAdvisorType() != null) {
-			return KeyValue.of(LowCardinalityKeyNames.ADVISOR_TYPE, context.getAdvisorType().name());
-		}
-		return ADVISOR_TYPE_NONE;
+		return KeyValue.of(LowCardinalityKeyNames.ADVISOR_TYPE, context.getAdvisorType().name());
 	}
 
 	protected KeyValue springAiKind() {
 		return KeyValue.of(AdvisorObservationDocumentation.LowCardinalityKeyNames.SPRING_AI_KIND,
-				CHAT_CLIENT_ADVISOR_SPRING_AI_KIND);
+				SpringAiKind.ADVISOR.value());
 	}
 
 	// ------------------------
@@ -94,14 +83,11 @@ public KeyValues getHighCardinalityKeyValues(AdvisorObservationContext context)
 	}
 
 	protected KeyValue advisorName(AdvisorObservationContext context) {
-		if (context.getAdvisorType() != null) {
-			return KeyValue.of(HighCardinalityKeyNames.ADVISOR_NAME, context.getAdvisorName());
-		}
-		return ADVISOR_NAME_NONE;
+		return KeyValue.of(HighCardinalityKeyNames.ADVISOR_NAME, context.getAdvisorName());
 	}
 
 	protected KeyValue advisorOrder(AdvisorObservationContext context) {
 		return KeyValue.of(HighCardinalityKeyNames.ADVISOR_ORDER, "" + context.getOrder());
 	}
 
-}
+}

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/observation/package-info.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2024 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.
+ */
+
+@NonNullApi
+@NonNullFields
+package org.springframework.ai.chat.client.advisor.observation;
+
+import org.springframework.lang.NonNullApi;
+import org.springframework.lang.NonNullFields;