Model observability for Anthropic

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>

models/spring-ai-anthropic/pom.xml

@@ -76,6 +76,12 @@
 			<scope>test</scope>
 		</dependency>
 
+		<dependency>
+			<groupId>io.micrometer</groupId>
+			<artifactId>micrometer-observation-test</artifactId>
+			<scope>test</scope>
+		</dependency>
+
 		<dependency>
 			<groupId>com.fasterxml.jackson.dataformat</groupId>
 			<artifactId>jackson-dataformat-xml</artifactId>

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java

@@ -23,6 +23,9 @@ import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.ai.anthropic.api.AnthropicApi;
@@ -39,11 +42,13 @@ import org.springframework.ai.chat.messages.ToolResponseMessage;
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
 import org.springframework.ai.chat.metadata.ChatResponseMetadata;
-import org.springframework.ai.chat.model.AbstractToolCallSupport;
-import org.springframework.ai.chat.model.ChatModel;
-import org.springframework.ai.chat.model.ChatResponse;
-import org.springframework.ai.chat.model.Generation;
+import org.springframework.ai.chat.model.*;
+import org.springframework.ai.chat.observation.ChatModelObservationContext;
+import org.springframework.ai.chat.observation.ChatModelObservationConvention;
+import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
+import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
 import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.model.ModelOptionsUtils;
 import org.springframework.ai.model.function.FunctionCallback;
@@ -64,12 +69,15 @@ import reactor.core.publisher.Mono;
  * @author Christian Tzolov
  * @author luocongqiu
  * @author Mariusz Bernacki
+ * @author Thomas Vitale
  * @since 1.0.0
  */
 public class AnthropicChatModel extends AbstractToolCallSupport implements ChatModel {
 
 	private static final Logger logger = LoggerFactory.getLogger(AnthropicChatModel.class);
 
+	private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
+
 	public static final String DEFAULT_MODEL_NAME = AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getValue();
 
 	public static final Integer DEFAULT_MAX_TOKENS = 500;
@@ -91,6 +99,16 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
 	 */
 	public final RetryTemplate retryTemplate;
 
+	/**
+	 * Observation registry used for instrumentation.
+	 */
+	private final ObservationRegistry observationRegistry;
+
+	/**
+	 * Conventions to use for generating observations.
+	 */
+	private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
+
 	/**
 	 * Construct a new {@link AnthropicChatModel} instance.
 	 * @param anthropicApi the lower-level API for the Anthropic service.
@@ -151,54 +169,108 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
 	public AnthropicChatModel(AnthropicApi anthropicApi, AnthropicChatOptions defaultOptions,
 			RetryTemplate retryTemplate, FunctionCallbackContext functionCallbackContext,
 			List<FunctionCallback> toolFunctionCallbacks) {
+		this(anthropicApi, defaultOptions, retryTemplate, functionCallbackContext, toolFunctionCallbacks,
+				ObservationRegistry.NOOP);
+	}
+
+	/**
+	 * Construct a new {@link AnthropicChatModel} instance.
+	 * @param anthropicApi the lower-level API for the Anthropic service.
+	 * @param defaultOptions the default options used for the chat completion requests.
+	 * @param retryTemplate the retry template used to retry the Anthropic API calls.
+	 * @param functionCallbackContext the function callback context used to store the
+	 * state of the function calls.
+	 * @param toolFunctionCallbacks the tool function callbacks used to handle the tool
+	 * calls.
+	 */
+	public AnthropicChatModel(AnthropicApi anthropicApi, AnthropicChatOptions defaultOptions,
+			RetryTemplate retryTemplate, FunctionCallbackContext functionCallbackContext,
+			List<FunctionCallback> toolFunctionCallbacks, ObservationRegistry observationRegistry) {
 
 		super(functionCallbackContext, defaultOptions, toolFunctionCallbacks);
 
 		Assert.notNull(anthropicApi, "AnthropicApi must not be null");
 		Assert.notNull(defaultOptions, "DefaultOptions must not be null");
 		Assert.notNull(retryTemplate, "RetryTemplate must not be null");
+		Assert.notNull(observationRegistry, "ObservationRegistry must not be null");
 
 		this.anthropicApi = anthropicApi;
 		this.defaultOptions = defaultOptions;
 		this.retryTemplate = retryTemplate;
+		this.observationRegistry = observationRegistry;
 	}
 
 	@Override
 	public ChatResponse call(Prompt prompt) {
-
 		ChatCompletionRequest request = createRequest(prompt, false);
 
-		ResponseEntity<ChatCompletionResponse> completionEntity = this.retryTemplate
-			.execute(ctx -> this.anthropicApi.chatCompletionEntity(request));
+		ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
+			.prompt(prompt)
+			.provider(AnthropicApi.PROVIDER_NAME)
+			.requestOptions(buildRequestOptions(request))
+			.build();
 
-		ChatResponse chatResponse = toChatResponse(completionEntity.getBody());
+		ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
+			.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					this.observationRegistry)
+			.observe(() -> {
 
-		if (this.isToolCall(chatResponse, Set.of("tool_use"))) {
-			var toolCallConversation = handleToolCalls(prompt, chatResponse);
+				ResponseEntity<ChatCompletionResponse> completionEntity = this.retryTemplate
+					.execute(ctx -> this.anthropicApi.chatCompletionEntity(request));
+
+				ChatResponse chatResponse = toChatResponse(completionEntity.getBody());
+
+				observationContext.setResponse(chatResponse);
+
+				return chatResponse;
+			});
+
+		if (response != null && this.isToolCall(response, Set.of("tool_use"))) {
+			var toolCallConversation = handleToolCalls(prompt, response);
 			return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
 		}
 
-		return chatResponse;
+		return response;
 	}
 
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
+		return Flux.deferContextual(contextView -> {
+			ChatCompletionRequest request = createRequest(prompt, true);
 
-		ChatCompletionRequest request = createRequest(prompt, true);
+			ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
+				.prompt(prompt)
+				.provider(AnthropicApi.PROVIDER_NAME)
+				.requestOptions(buildRequestOptions(request))
+				.build();
 
-		Flux<ChatCompletionResponse> response = this.retryTemplate
-			.execute(ctx -> this.anthropicApi.chatCompletionStream(request));
+			Observation observation = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION.observation(
+					this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
+					this.observationRegistry);
 
-		return response.switchMap(chatCompletionResponse -> {
+			observation.parentObservation(contextView.getOrDefault(ObservationThreadLocalAccessor.KEY, null)).start();
 
-			ChatResponse chatResponse = toChatResponse(chatCompletionResponse);
+			Flux<ChatCompletionResponse> response = this.anthropicApi.chatCompletionStream(request);
 
-			if (this.isToolCall(chatResponse, Set.of("tool_use"))) {
-				var toolCallConversation = handleToolCalls(prompt, chatResponse);
-				return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
-			}
+			// @formatter:off
+			Flux<ChatResponse> chatResponseFlux = response.switchMap(chatCompletionResponse -> {
+				ChatResponse chatResponse = toChatResponse(chatCompletionResponse);
 
-			return Mono.just(chatResponse);
+				if (this.isToolCall(chatResponse, Set.of("tool_use"))) {
+					var toolCallConversation = handleToolCalls(prompt, chatResponse);
+					return this.stream(new Prompt(toolCallConversation, prompt.getOptions()));
+				}
+
+				return Mono.just(chatResponse);
+			})
+			.doOnError(observation::error)
+			.doFinally(s -> {
+				observation.stop();
+			})
+			.contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
+			// @formatter:on
+
+			return new MessageAggregator().aggregate(chatResponseFlux, observationContext::setResponse);
 		});
 	}
 
@@ -366,9 +438,29 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM
 		}).toList();
 	}
 
+	private ChatOptions buildRequestOptions(AnthropicApi.ChatCompletionRequest request) {
+		return ChatOptionsBuilder.builder()
+			.withModel(request.model())
+			.withMaxTokens(request.maxTokens())
+			.withStopSequences(request.stopSequences())
+			.withTemperature(request.temperature())
+			.withTopK(request.topK())
+			.withTopP(request.topP())
+			.build();
+	}
+
 	@Override
 	public ChatOptions getDefaultOptions() {
 		return AnthropicChatOptions.fromOptions(this.defaultOptions);
 	}
 
+	/**
+	 * Use the provided convention for reporting observation data
+	 * @param observationConvention The provided convention
+	 */
+	public void setObservationConvention(ChatModelObservationConvention observationConvention) {
+		Assert.notNull(observationConvention, "observationConvention cannot be null");
+		this.observationConvention = observationConvention;
+	}
+
 }

models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/api/AnthropicApi.java

@@ -23,11 +23,10 @@ import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.springframework.ai.anthropic.api.StreamHelper.ChatCompletionResponseBuilder;
 import org.springframework.ai.model.ChatModelDescription;
 import org.springframework.ai.model.ModelOptionsUtils;
+import org.springframework.ai.observation.conventions.AiProvider;
 import org.springframework.ai.retry.RetryUtils;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatusCode;
@@ -51,11 +50,12 @@ import reactor.core.publisher.Mono;
 /**
  * @author Christian Tzolov
  * @author Mariusz Bernacki
+ * @author Thomas Vitale
  * @since 1.0.0
  */
 public class AnthropicApi {
 
-	private static final Logger logger = LoggerFactory.getLogger(AnthropicApi.class);
+	public static final String PROVIDER_NAME = AiProvider.ANTHROPIC.value();
 
 	private static final String HEADER_X_API_KEY = "x-api-key";
 

models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/AnthropicChatModelObservationIT.java

@@ -0,0 +1,179 @@
+/*
+ * 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.
+ */
+package org.springframework.ai.anthropic;
+
+import io.micrometer.common.KeyValue;
+import io.micrometer.observation.tck.TestObservationRegistry;
+import io.micrometer.observation.tck.TestObservationRegistryAssert;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.springframework.ai.anthropic.api.AnthropicApi;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.model.function.FunctionCallbackContext;
+import org.springframework.ai.observation.conventions.AiOperationType;
+import org.springframework.ai.observation.conventions.AiProvider;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+import org.springframework.retry.support.RetryTemplate;
+import reactor.core.publisher.Flux;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.HighCardinalityKeyNames;
+import static org.springframework.ai.chat.observation.ChatModelObservationDocumentation.LowCardinalityKeyNames;
+
+/**
+ * Integration tests for observation instrumentation in {@link AnthropicChatModel}.
+ *
+ * @author Thomas Vitale
+ */
+@SpringBootTest(classes = AnthropicChatModelObservationIT.Config.class,
+		properties = "spring.ai.retry.on-http-codes=429")
+@EnabledIfEnvironmentVariable(named = "ANTHROPIC_API_KEY", matches = ".+")
+public class AnthropicChatModelObservationIT {
+
+	@Autowired
+	TestObservationRegistry observationRegistry;
+
+	@Autowired
+	AnthropicChatModel chatModel;
+
+	@BeforeEach
+	void beforeEach() {
+		observationRegistry.clear();
+	}
+
+	@Test
+	void observationForChatOperation() {
+		var options = AnthropicChatOptions.builder()
+			.withModel(AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getValue())
+			.withMaxTokens(2048)
+			.withStopSequences(List.of("this-is-the-end"))
+			.withTemperature(0.7f)
+			.withTopK(1)
+			.withTopP(1f)
+			.build();
+
+		Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
+
+		ChatResponse chatResponse = chatModel.call(prompt);
+		assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty();
+
+		ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
+		assertThat(responseMetadata).isNotNull();
+
+		validate(responseMetadata, "[\"end_turn\"]");
+	}
+
+	@Test
+	void observationForStreamingChatOperation() {
+		var options = AnthropicChatOptions.builder()
+			.withModel(AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getValue())
+			.withMaxTokens(2048)
+			.withStopSequences(List.of("this-is-the-end"))
+			.withTemperature(0.7f)
+			.withTopK(1)
+			.withTopP(1f)
+			.build();
+
+		Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
+
+		Flux<ChatResponse> chatResponseFlux = chatModel.stream(prompt);
+
+		List<ChatResponse> responses = chatResponseFlux.collectList().block();
+		assertThat(responses).isNotEmpty();
+		assertThat(responses).hasSizeGreaterThan(3);
+
+		String aggregatedResponse = responses.subList(0, responses.size() - 1)
+			.stream()
+			.filter(r -> r.getResult() != null)
+			.map(r -> r.getResult().getOutput().getContent())
+			.collect(Collectors.joining());
+		assertThat(aggregatedResponse).isNotEmpty();
+
+		ChatResponse lastChatResponse = responses.get(responses.size() - 1);
+
+		ChatResponseMetadata responseMetadata = lastChatResponse.getMetadata();
+		assertThat(responseMetadata).isNotNull();
+
+		validate(responseMetadata, KeyValue.NONE_VALUE);
+	}
+
+	private void validate(ChatResponseMetadata responseMetadata, String finishReasons) {
+		TestObservationRegistryAssert.assertThat(observationRegistry)
+			.doesNotHaveAnyRemainingCurrentObservation()
+			.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)
+			.that()
+			.hasContextualNameEqualTo("chat " + AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getValue())
+			.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
+					AiOperationType.CHAT.value())
+			.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.ANTHROPIC.value())
+			.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(),
+					AnthropicApi.ChatModel.CLAUDE_3_5_SONNET.getValue())
+			.hasLowCardinalityKeyValue(LowCardinalityKeyNames.RESPONSE_MODEL.asString(), responseMetadata.getModel())
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString(),
+					KeyValue.NONE_VALUE)
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString(),
+					KeyValue.NONE_VALUE)
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(),
+					"[\"this-is-the-end\"]")
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7")
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_K.asString(), "1")
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0")
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_ID.asString(), responseMetadata.getId())
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(), finishReasons)
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
+					String.valueOf(responseMetadata.getUsage().getPromptTokens()))
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(),
+					String.valueOf(responseMetadata.getUsage().getGenerationTokens()))
+			.hasHighCardinalityKeyValue(HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
+					String.valueOf(responseMetadata.getUsage().getTotalTokens()))
+			.hasBeenStarted()
+			.hasBeenStopped();
+	}
+
+	@SpringBootConfiguration
+	static class Config {
+
+		@Bean
+		public TestObservationRegistry observationRegistry() {
+			return TestObservationRegistry.create();
+		}
+
+		@Bean
+		public AnthropicApi anthropicApi() {
+			return new AnthropicApi(System.getenv("ANTHROPIC_API_KEY"));
+		}
+
+		@Bean
+		public AnthropicChatModel anthropicChatModel(AnthropicApi anthropicApi,
+				TestObservationRegistry observationRegistry) {
+			return new AnthropicChatModel(anthropicApi, AnthropicChatOptions.builder().build(),
+					RetryTemplate.defaultInstance(), new FunctionCallbackContext(), List.of(), observationRegistry);
+		}
+
+	}
+
+}