Skip to content

Commit c6bb075

Browse files
sobychackosobychacko
committed
Add Micrometer observability for Azure OpenAI chat model imperative calls
- Integrate ObservationRegistry in AzureOpenAiChatModel - Implement observation context and conventions - Add integration test for chat model observations - Update AiProvider enum with AZURE_OPENAI entry - Support Micrometer observability for imperative call()
1 parent 05292ac commit c6bb075

File tree

4 files changed

+251
-25
lines changed

4 files changed

+251
-25
lines changed

‎models/spring-ai-azure-openai/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,12 @@
5858
<artifactId>spring-boot-starter-test</artifactId>
5959
<scope>test</scope>
6060
</dependency>
61+
62+
<dependency>
63+
<groupId>io.micrometer</groupId>
64+
<artifactId>micrometer-observation-test</artifactId>
65+
<scope>test</scope>
66+
</dependency>
6167
</dependencies>
6268

6369
</project>

‎models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java‎

Lines changed: 84 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,16 @@
1616

1717
package org.springframework.ai.azure.openai;
1818

19-
import com.azure.ai.openai.OpenAIAsyncClient;
20-
import com.azure.ai.openai.OpenAIClient;
21-
import com.azure.ai.openai.OpenAIClientBuilder;
22-
import com.azure.ai.openai.models.*;
23-
import com.azure.core.util.BinaryData;
19+
import java.util.ArrayList;
20+
import java.util.Base64;
21+
import java.util.Collections;
22+
import java.util.HashSet;
23+
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Optional;
26+
import java.util.Set;
27+
import java.util.concurrent.atomic.AtomicBoolean;
28+
2429
import org.springframework.ai.azure.openai.metadata.AzureOpenAiUsage;
2530
import org.springframework.ai.chat.messages.AssistantMessage;
2631
import org.springframework.ai.chat.messages.Message;
@@ -36,27 +41,51 @@
3641
import org.springframework.ai.chat.model.ChatModel;
3742
import org.springframework.ai.chat.model.ChatResponse;
3843
import org.springframework.ai.chat.model.Generation;
44+
import org.springframework.ai.chat.observation.ChatModelObservationContext;
45+
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
46+
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
47+
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
3948
import org.springframework.ai.chat.prompt.ChatOptions;
4049
import org.springframework.ai.chat.prompt.Prompt;
4150
import org.springframework.ai.model.Media;
4251
import org.springframework.ai.model.ModelOptionsUtils;
4352
import org.springframework.ai.model.function.FunctionCallback;
4453
import org.springframework.ai.model.function.FunctionCallbackContext;
54+
import org.springframework.ai.observation.conventions.AiProvider;
4555
import org.springframework.util.Assert;
4656
import org.springframework.util.CollectionUtils;
57+
58+
import com.azure.ai.openai.OpenAIAsyncClient;
59+
import com.azure.ai.openai.OpenAIClient;
60+
import com.azure.ai.openai.OpenAIClientBuilder;
61+
import com.azure.ai.openai.models.ChatChoice;
62+
import com.azure.ai.openai.models.ChatCompletions;
63+
import com.azure.ai.openai.models.ChatCompletionsFunctionToolCall;
64+
import com.azure.ai.openai.models.ChatCompletionsFunctionToolDefinition;
65+
import com.azure.ai.openai.models.ChatCompletionsJsonResponseFormat;
66+
import com.azure.ai.openai.models.ChatCompletionsOptions;
67+
import com.azure.ai.openai.models.ChatCompletionsResponseFormat;
68+
import com.azure.ai.openai.models.ChatCompletionsTextResponseFormat;
69+
import com.azure.ai.openai.models.ChatCompletionsToolCall;
70+
import com.azure.ai.openai.models.ChatCompletionsToolDefinition;
71+
import com.azure.ai.openai.models.ChatMessageContentItem;
72+
import com.azure.ai.openai.models.ChatMessageImageContentItem;
73+
import com.azure.ai.openai.models.ChatMessageImageUrl;
74+
import com.azure.ai.openai.models.ChatMessageTextContentItem;
75+
import com.azure.ai.openai.models.ChatRequestAssistantMessage;
76+
import com.azure.ai.openai.models.ChatRequestMessage;
77+
import com.azure.ai.openai.models.ChatRequestSystemMessage;
78+
import com.azure.ai.openai.models.ChatRequestToolMessage;
79+
import com.azure.ai.openai.models.ChatRequestUserMessage;
80+
import com.azure.ai.openai.models.CompletionsFinishReason;
81+
import com.azure.ai.openai.models.ContentFilterResultsForPrompt;
82+
import com.azure.ai.openai.models.FunctionCall;
83+
import com.azure.ai.openai.models.FunctionDefinition;
84+
import com.azure.core.util.BinaryData;
85+
import io.micrometer.observation.ObservationRegistry;
4786
import reactor.core.publisher.Flux;
4887
import reactor.core.publisher.Mono;
4988

50-
import java.util.ArrayList;
51-
import java.util.Base64;
52-
import java.util.Collections;
53-
import java.util.HashSet;
54-
import java.util.List;
55-
import java.util.Map;
56-
import java.util.Optional;
57-
import java.util.Set;
58-
import java.util.concurrent.atomic.AtomicBoolean;
59-
6089
/**
6190
* {@link ChatModel} implementation for {@literal Microsoft Azure AI} backed by
6291
* {@link OpenAIClient}.
@@ -81,6 +110,8 @@ public class AzureOpenAiChatModel extends AbstractToolCallSupport implements Cha
81110

82111
private static final Double DEFAULT_TEMPERATURE = 0.7;
83112

113+
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
114+
84115
/**
85116
* The {@link OpenAIClient} used to interact with the Azure OpenAI service.
86117
*/
@@ -96,8 +127,18 @@ public class AzureOpenAiChatModel extends AbstractToolCallSupport implements Cha
96127
*/
97128
private final AzureOpenAiChatOptions defaultOptions;
98129

99-
public AzureOpenAiChatModel(OpenAIClientBuilder microsoftOpenAiClient) {
100-
this(microsoftOpenAiClient,
130+
/**
131+
* Observation registry used for instrumentation.
132+
*/
133+
private final ObservationRegistry observationRegistry;
134+
135+
/**
136+
* Conventions to use for generating observations.
137+
*/
138+
private ChatModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
139+
140+
public AzureOpenAiChatModel(OpenAIClientBuilder openAIClientBuilder) {
141+
this(openAIClientBuilder,
101142
AzureOpenAiChatOptions.builder()
102143
.withDeploymentName(DEFAULT_DEPLOYMENT_NAME)
103144
.withTemperature(DEFAULT_TEMPERATURE)
@@ -115,12 +156,19 @@ public AzureOpenAiChatModel(OpenAIClientBuilder openAIClientBuilder, AzureOpenAi
115156

116157
public AzureOpenAiChatModel(OpenAIClientBuilder openAIClientBuilder, AzureOpenAiChatOptions options,
117158
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks) {
159+
this(openAIClientBuilder, options, functionCallbackContext, List.of(), ObservationRegistry.NOOP);
160+
}
161+
162+
public AzureOpenAiChatModel(OpenAIClientBuilder openAIClientBuilder, AzureOpenAiChatOptions options,
163+
FunctionCallbackContext functionCallbackContext, List<FunctionCallback> toolFunctionCallbacks,
164+
ObservationRegistry observationRegistry) {
118165
super(functionCallbackContext, options, toolFunctionCallbacks);
119166
Assert.notNull(openAIClientBuilder, "com.azure.ai.openai.OpenAIClient must not be null");
120167
Assert.notNull(options, "AzureOpenAiChatOptions must not be null");
121168
this.openAIClient = openAIClientBuilder.buildClient();
122169
this.openAIAsyncClient = openAIClientBuilder.buildAsyncClient();
123170
this.defaultOptions = options;
171+
this.observationRegistry = observationRegistry;
124172
}
125173

126174
public AzureOpenAiChatOptions getDefaultOptions() {
@@ -130,22 +178,34 @@ public AzureOpenAiChatOptions getDefaultOptions() {
130178
@Override
131179
public ChatResponse call(Prompt prompt) {
132180

133-
ChatCompletionsOptions options = toAzureChatCompletionsOptions(prompt);
134-
options.setStream(false);
181+
ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
182+
.prompt(prompt)
183+
.provider(AiProvider.AZURE_OPENAI.value())
184+
.requestOptions(prompt.getOptions() != null ? prompt.getOptions() : this.defaultOptions)
185+
.build();
135186

136-
ChatCompletions chatCompletions = this.openAIClient.getChatCompletions(options.getModel(), options);
187+
ChatResponse response = ChatModelObservationDocumentation.CHAT_MODEL_OPERATION
188+
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
189+
this.observationRegistry)
190+
.observe(() -> {
191+
ChatCompletionsOptions options = toAzureChatCompletionsOptions(prompt);
192+
options.setStream(false);
137193

138-
ChatResponse chatResponse = toChatResponse(chatCompletions);
194+
ChatCompletions chatCompletions = this.openAIClient.getChatCompletions(options.getModel(), options);
195+
ChatResponse chatResponse = toChatResponse(chatCompletions);
196+
observationContext.setResponse(chatResponse);
197+
return chatResponse;
198+
});
139199

140200
if (!isProxyToolCalls(prompt, this.defaultOptions)
141-
&& isToolCall(chatResponse, Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
142-
var toolCallConversation = handleToolCalls(prompt, chatResponse);
201+
&& isToolCall(response, Set.of(String.valueOf(CompletionsFinishReason.TOOL_CALLS).toLowerCase()))) {
202+
var toolCallConversation = handleToolCalls(prompt, response);
143203
// Recursively call the call method with the tool call message
144204
// conversation that contains the call responses.
145205
return this.call(new Prompt(toolCallConversation, prompt.getOptions()));
146206
}
147207

148-
return chatResponse;
208+
return response;
149209
}
150210

151211
@Override

‎models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatModelObservationIT.java‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.ai.azure.openai;
17+
18+
import static com.azure.core.http.policy.HttpLogDetailLevel.BODY_AND_HEADERS;
19+
import static org.assertj.core.api.Assertions.assertThat;
20+
21+
import java.util.List;
22+
23+
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
25+
26+
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
27+
import org.springframework.ai.chat.model.ChatResponse;
28+
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
29+
import org.springframework.ai.chat.observation.DefaultChatModelObservationConvention;
30+
import org.springframework.ai.chat.prompt.Prompt;
31+
import org.springframework.ai.observation.conventions.AiOperationType;
32+
import org.springframework.ai.observation.conventions.AiProvider;
33+
import org.springframework.beans.factory.annotation.Autowired;
34+
import org.springframework.boot.SpringBootConfiguration;
35+
import org.springframework.boot.test.context.SpringBootTest;
36+
import org.springframework.context.annotation.Bean;
37+
38+
import com.azure.ai.openai.OpenAIClientBuilder;
39+
import com.azure.ai.openai.OpenAIServiceVersion;
40+
import com.azure.core.credential.AzureKeyCredential;
41+
import com.azure.core.http.policy.HttpLogOptions;
42+
import io.micrometer.common.KeyValue;
43+
import io.micrometer.observation.tck.TestObservationRegistry;
44+
import io.micrometer.observation.tck.TestObservationRegistryAssert;
45+
46+
/**
47+
* @author Soby Chacko
48+
*/
49+
@SpringBootTest(classes = AzureOpenAiChatModelObservationIT.TestConfiguration.class)
50+
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_API_KEY", matches = ".+")
51+
@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_ENDPOINT", matches = ".+")
52+
class AzureOpenAiChatModelObservationIT {
53+
54+
@Autowired
55+
private AzureOpenAiChatModel chatModel;
56+
57+
@Autowired
58+
TestObservationRegistry observationRegistry;
59+
60+
@Test
61+
void observationForImperativeChatOperation() {
62+
63+
var options = AzureOpenAiChatOptions.builder()
64+
.withFrequencyPenalty(0.0)
65+
.withMaxTokens(2048)
66+
.withPresencePenalty(0.0)
67+
.withStop(List.of("this-is-the-end"))
68+
.withTemperature(0.7)
69+
.withTopP(1.0)
70+
.build();
71+
72+
Prompt prompt = new Prompt("Why does a raven look like a desk?", options);
73+
74+
ChatResponse chatResponse = chatModel.call(prompt);
75+
assertThat(chatResponse.getResult().getOutput().getContent()).isNotEmpty();
76+
77+
ChatResponseMetadata responseMetadata = chatResponse.getMetadata();
78+
assertThat(responseMetadata).isNotNull();
79+
80+
validate(responseMetadata);
81+
}
82+
83+
private void validate(ChatResponseMetadata responseMetadata) {
84+
TestObservationRegistryAssert.assertThat(observationRegistry)
85+
.doesNotHaveAnyRemainingCurrentObservation()
86+
.hasObservationWithNameEqualTo(DefaultChatModelObservationConvention.DEFAULT_NAME)
87+
.that()
88+
.hasLowCardinalityKeyValue(
89+
ChatModelObservationDocumentation.LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
90+
AiOperationType.CHAT.value())
91+
.hasLowCardinalityKeyValue(ChatModelObservationDocumentation.LowCardinalityKeyNames.AI_PROVIDER.asString(),
92+
AiProvider.AZURE_OPENAI.value())
93+
.hasLowCardinalityKeyValue(
94+
ChatModelObservationDocumentation.LowCardinalityKeyNames.RESPONSE_MODEL.asString(),
95+
responseMetadata.getModel())
96+
.hasHighCardinalityKeyValue(
97+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_FREQUENCY_PENALTY.asString(),
98+
"0.0")
99+
.hasHighCardinalityKeyValue(
100+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
101+
.hasHighCardinalityKeyValue(
102+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_PRESENCE_PENALTY.asString(),
103+
"0.0")
104+
.hasHighCardinalityKeyValue(
105+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_STOP_SEQUENCES.asString(),
106+
"[\"this-is-the-end\"]")
107+
.hasHighCardinalityKeyValue(
108+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TEMPERATURE.asString(), "0.7")
109+
.hasHighCardinalityKeyValue(
110+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TOP_K.asString(),
111+
KeyValue.NONE_VALUE)
112+
.hasHighCardinalityKeyValue(
113+
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_TOP_P.asString(), "1.0")
114+
.hasHighCardinalityKeyValue(
115+
ChatModelObservationDocumentation.HighCardinalityKeyNames.RESPONSE_ID.asString(),
116+
responseMetadata.getId())
117+
.hasHighCardinalityKeyValue(
118+
ChatModelObservationDocumentation.HighCardinalityKeyNames.RESPONSE_FINISH_REASONS.asString(),
119+
"[\"stop\"]")
120+
.hasHighCardinalityKeyValue(
121+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_INPUT_TOKENS.asString(),
122+
String.valueOf(responseMetadata.getUsage().getPromptTokens()))
123+
.hasHighCardinalityKeyValue(
124+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_OUTPUT_TOKENS.asString(),
125+
String.valueOf(responseMetadata.getUsage().getGenerationTokens()))
126+
.hasHighCardinalityKeyValue(
127+
ChatModelObservationDocumentation.HighCardinalityKeyNames.USAGE_TOTAL_TOKENS.asString(),
128+
String.valueOf(responseMetadata.getUsage().getTotalTokens()))
129+
.hasBeenStarted()
130+
.hasBeenStopped();
131+
}
132+
133+
@SpringBootConfiguration
134+
public static class TestConfiguration {
135+
136+
@Bean
137+
public TestObservationRegistry observationRegistry() {
138+
return TestObservationRegistry.create();
139+
}
140+
141+
@Bean
142+
public OpenAIClientBuilder openAIClient() {
143+
return new OpenAIClientBuilder().credential(new AzureKeyCredential(System.getenv("AZURE_OPENAI_API_KEY")))
144+
.endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
145+
.serviceVersion(OpenAIServiceVersion.V2024_02_15_PREVIEW)
146+
.httpLogOptions(new HttpLogOptions().setLogLevel(BODY_AND_HEADERS));
147+
}
148+
149+
@Bean
150+
public AzureOpenAiChatModel azureOpenAiChatModel(OpenAIClientBuilder openAIClientBuilder,
151+
TestObservationRegistry observationRegistry) {
152+
return new AzureOpenAiChatModel(openAIClientBuilder,
153+
AzureOpenAiChatOptions.builder().withDeploymentName("gpt-4o").withMaxTokens(1000).build(), null,
154+
List.of(), observationRegistry);
155+
}
156+
157+
}
158+
159+
}

‎spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiProvider.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ public enum AiProvider {
3535
OPENAI("openai"),
3636
SPRING_AI("spring_ai"),
3737
VERTEX_AI("vertex_ai"),
38-
OCI_GENAI("oci_genai");
38+
OCI_GENAI("oci_genai"),
39+
AZURE_OPENAI("azure-openai");
3940

4041
private final String value;
4142

0 commit comments

Comments
 (0)