Skip to content

Commit 08d2ceb

Browse files
mxsl-grmxsl-gr
committed
feat: add observability to QianFan image model
1 parent cddb00a commit 08d2ceb

File tree

3 files changed

+218
-45
lines changed

3 files changed

+218
-45
lines changed

‎models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanImageModel.java‎

Lines changed: 98 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package org.springframework.ai.qianfan;
1717

18+
import io.micrometer.observation.ObservationRegistry;
1819
import org.slf4j.Logger;
1920
import org.slf4j.LoggerFactory;
2021
import org.springframework.ai.image.Image;
@@ -23,10 +24,16 @@
2324
import org.springframework.ai.image.ImageOptions;
2425
import org.springframework.ai.image.ImagePrompt;
2526
import org.springframework.ai.image.ImageResponse;
27+
import org.springframework.ai.image.observation.DefaultImageModelObservationConvention;
28+
import org.springframework.ai.image.observation.ImageModelObservationContext;
29+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
30+
import org.springframework.ai.image.observation.ImageModelObservationDocumentation;
2631
import org.springframework.ai.model.ModelOptionsUtils;
32+
import org.springframework.ai.qianfan.api.QianFanConstants;
2733
import org.springframework.ai.qianfan.api.QianFanImageApi;
2834
import org.springframework.ai.retry.RetryUtils;
2935
import org.springframework.http.ResponseEntity;
36+
import org.springframework.lang.Nullable;
3037
import org.springframework.retry.support.RetryTemplate;
3138
import org.springframework.util.Assert;
3239

@@ -43,6 +50,8 @@ public class QianFanImageModel implements ImageModel {
4350

4451
private final static Logger logger = LoggerFactory.getLogger(QianFanImageModel.class);
4552

53+
private static final ImageModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultImageModelObservationConvention();
54+
4655
/**
4756
* The default options used for the image completion requests.
4857
*/
@@ -58,6 +67,16 @@ public class QianFanImageModel implements ImageModel {
5867
*/
5968
private final QianFanImageApi qianFanImageApi;
6069

70+
/**
71+
* Observation registry used for instrumentation.
72+
*/
73+
private final ObservationRegistry observationRegistry;
74+
75+
/**
76+
* Conventions to use for generating observations.
77+
*/
78+
private ImageModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
79+
6180
/**
6281
* Creates an instance of the QianFanImageModel.
6382
* @param qianFanImageApi The QianFanImageApi instance to be used for interacting with
@@ -69,48 +88,85 @@ public QianFanImageModel(QianFanImageApi qianFanImageApi) {
6988
}
7089

7190
/**
72-
* Initializes a new instance of the QianFanImageModel.
91+
* Creates an instance of the QianFanImageModel.
92+
* @param qianFanImageApi The QianFanImageApi instance to be used for interacting with
93+
* the QianFan Image API.
94+
* @param options The QianFanImageOptions to configure the image model.
95+
* @throws IllegalArgumentException if qianFanImageApi is null
96+
*/
97+
public QianFanImageModel(QianFanImageApi qianFanImageApi, QianFanImageOptions options) {
98+
this(qianFanImageApi, options, RetryUtils.DEFAULT_RETRY_TEMPLATE);
99+
}
100+
101+
/**
102+
* Creates an instance of the QianFanImageModel.
73103
* @param qianFanImageApi The QianFanImageApi instance to be used for interacting with
74104
* the QianFan Image API.
75105
* @param options The QianFanImageOptions to configure the image model.
76106
* @param retryTemplate The retry template.
107+
* @throws IllegalArgumentException if qianFanImageApi is null
77108
*/
78109
public QianFanImageModel(QianFanImageApi qianFanImageApi, QianFanImageOptions options,
79110
RetryTemplate retryTemplate) {
111+
this(qianFanImageApi, options, retryTemplate, ObservationRegistry.NOOP);
112+
}
113+
114+
/**
115+
* Initializes a new instance of the QianFanImageModel.
116+
* @param qianFanImageApi The QianFanImageApi instance to be used for interacting with
117+
* the QianFan Image API.
118+
* @param options The QianFanImageOptions to configure the image model.
119+
* @param retryTemplate The retry template.
120+
* @param observationRegistry The ObservationRegistry used for instrumentation.
121+
*/
122+
public QianFanImageModel(QianFanImageApi qianFanImageApi, QianFanImageOptions options, RetryTemplate retryTemplate,
123+
ObservationRegistry observationRegistry) {
80124
Assert.notNull(qianFanImageApi, "QianFanImageApi must not be null");
81125
Assert.notNull(options, "options must not be null");
82126
Assert.notNull(retryTemplate, "retryTemplate must not be null");
127+
Assert.notNull(observationRegistry, "observationRegistry must not be null");
83128
this.qianFanImageApi = qianFanImageApi;
84129
this.defaultOptions = options;
85130
this.retryTemplate = retryTemplate;
131+
this.observationRegistry = observationRegistry;
86132
}
87133

88134
@Override
89135
public ImageResponse call(ImagePrompt imagePrompt) {
90-
return this.retryTemplate.execute(ctx -> {
136+
QianFanImageOptions requestImageOptions = mergeOptions(imagePrompt.getOptions(), this.defaultOptions);
91137

92-
String instructions = imagePrompt.getInstructions().get(0).getText();
138+
QianFanImageApi.QianFanImageRequest imageRequest = createRequest(imagePrompt, requestImageOptions);
93139

94-
QianFanImageApi.QianFanImageRequest imageRequest = new QianFanImageApi.QianFanImageRequest(instructions,
95-
QianFanImageApi.DEFAULT_IMAGE_MODEL);
140+
var observationContext = ImageModelObservationContext.builder()
141+
.imagePrompt(imagePrompt)
142+
.provider(QianFanConstants.PROVIDER_NAME)
143+
.requestOptions(requestImageOptions)
144+
.build();
96145

97-
if (this.defaultOptions != null) {
98-
imageRequest = ModelOptionsUtils.merge(this.defaultOptions, imageRequest,
99-
QianFanImageApi.QianFanImageRequest.class);
100-
}
146+
return ImageModelObservationDocumentation.IMAGE_MODEL_OPERATION
147+
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
148+
this.observationRegistry)
149+
.observe(() -> {
101150

102-
if (imagePrompt.getOptions() != null) {
103-
imageRequest = ModelOptionsUtils.merge(toQianFanImageOptions(imagePrompt.getOptions()), imageRequest,
104-
QianFanImageApi.QianFanImageRequest.class);
105-
}
151+
ResponseEntity<QianFanImageApi.QianFanImageResponse> imageResponseEntity = this.retryTemplate
152+
.execute(ctx -> this.qianFanImageApi.createImage(imageRequest));
106153

107-
// Make the request
108-
ResponseEntity<QianFanImageApi.QianFanImageResponse> imageResponseEntity = this.qianFanImageApi
109-
.createImage(imageRequest);
154+
ImageResponse imageResponse = convertResponse(imageResponseEntity, imageRequest);
110155

111-
// Convert to org.springframework.ai.model derived ImageResponse data type
112-
return convertResponse(imageResponseEntity, imageRequest);
113-
});
156+
observationContext.setResponse(imageResponse);
157+
158+
return imageResponse;
159+
});
160+
}
161+
162+
private QianFanImageApi.QianFanImageRequest createRequest(ImagePrompt imagePrompt,
163+
QianFanImageOptions requestImageOptions) {
164+
String instructions = imagePrompt.getInstructions().get(0).getText();
165+
166+
QianFanImageApi.QianFanImageRequest imageRequest = new QianFanImageApi.QianFanImageRequest(instructions,
167+
QianFanImageApi.DEFAULT_IMAGE_MODEL);
168+
169+
return ModelOptionsUtils.merge(requestImageOptions, imageRequest, QianFanImageApi.QianFanImageRequest.class);
114170
}
115171

116172
private ImageResponse convertResponse(ResponseEntity<QianFanImageApi.QianFanImageResponse> imageResponseEntity,
@@ -132,33 +188,32 @@ private ImageResponse convertResponse(ResponseEntity<QianFanImageApi.QianFanImag
132188
/**
133189
* Convert the {@link ImageOptions} into {@link QianFanImageOptions}.
134190
* @param runtimeImageOptions the image options to use.
191+
* @param defaultOptions the default options.
135192
* @return the converted {@link QianFanImageOptions}.
136193
*/
137-
private QianFanImageOptions toQianFanImageOptions(ImageOptions runtimeImageOptions) {
138-
QianFanImageOptions.Builder qianFanImageOptionsBuilder = QianFanImageOptions.builder();
139-
if (runtimeImageOptions != null) {
140-
if (runtimeImageOptions.getN() != null) {
141-
qianFanImageOptionsBuilder.withN(runtimeImageOptions.getN());
142-
}
143-
if (runtimeImageOptions.getModel() != null) {
144-
qianFanImageOptionsBuilder.withModel(runtimeImageOptions.getModel());
145-
}
146-
if (runtimeImageOptions.getWidth() != null) {
147-
qianFanImageOptionsBuilder.withWidth(runtimeImageOptions.getWidth());
148-
}
149-
if (runtimeImageOptions.getHeight() != null) {
150-
qianFanImageOptionsBuilder.withHeight(runtimeImageOptions.getHeight());
151-
}
152-
if (runtimeImageOptions instanceof QianFanImageOptions runtimeQianFanImageOptions) {
153-
if (runtimeQianFanImageOptions.getStyle() != null) {
154-
qianFanImageOptionsBuilder.withStyle(runtimeQianFanImageOptions.getStyle());
155-
}
156-
if (runtimeQianFanImageOptions.getUser() != null) {
157-
qianFanImageOptionsBuilder.withUser(runtimeQianFanImageOptions.getUser());
158-
}
159-
}
194+
private QianFanImageOptions mergeOptions(@Nullable ImageOptions runtimeImageOptions,
195+
QianFanImageOptions defaultOptions) {
196+
var runtimeOptionsForProvider = ModelOptionsUtils.copyToTarget(runtimeImageOptions, ImageOptions.class,
197+
QianFanImageOptions.class);
198+
199+
if (runtimeOptionsForProvider == null) {
200+
return defaultOptions;
160201
}
161-
return qianFanImageOptionsBuilder.build();
202+
203+
return QianFanImageOptions.builder()
204+
.withModel(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getModel(), defaultOptions.getModel()))
205+
.withN(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getN(), defaultOptions.getN()))
206+
.withModel(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getModel(), defaultOptions.getModel()))
207+
.withWidth(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getWidth(), defaultOptions.getWidth()))
208+
.withHeight(
209+
ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getHeight(), defaultOptions.getHeight()))
210+
.withStyle(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getStyle(), defaultOptions.getStyle()))
211+
.withUser(ModelOptionsUtils.mergeOption(runtimeOptionsForProvider.getUser(), defaultOptions.getUser()))
212+
.build();
213+
}
214+
215+
public void setObservationConvention(ImageModelObservationConvention observationConvention) {
216+
this.observationConvention = observationConvention;
162217
}
163218

164219
}

‎models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/image/QianFanImageModelObservationIT.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* Copyright 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.qianfan.image;
17+
18+
import io.micrometer.observation.tck.TestObservationRegistry;
19+
import io.micrometer.observation.tck.TestObservationRegistryAssert;
20+
import org.junit.jupiter.api.Test;
21+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
22+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariables;
23+
import org.springframework.ai.image.ImagePrompt;
24+
import org.springframework.ai.image.ImageResponse;
25+
import org.springframework.ai.image.observation.DefaultImageModelObservationConvention;
26+
import org.springframework.ai.observation.conventions.AiOperationType;
27+
import org.springframework.ai.observation.conventions.AiProvider;
28+
import org.springframework.ai.qianfan.QianFanImageModel;
29+
import org.springframework.ai.qianfan.QianFanImageOptions;
30+
import org.springframework.ai.qianfan.api.QianFanImageApi;
31+
import org.springframework.beans.factory.annotation.Autowired;
32+
import org.springframework.boot.SpringBootConfiguration;
33+
import org.springframework.boot.test.context.SpringBootTest;
34+
import org.springframework.context.annotation.Bean;
35+
import org.springframework.retry.support.RetryTemplate;
36+
37+
import static org.assertj.core.api.Assertions.assertThat;
38+
import static org.springframework.ai.image.observation.ImageModelObservationDocumentation.HighCardinalityKeyNames;
39+
import static org.springframework.ai.image.observation.ImageModelObservationDocumentation.LowCardinalityKeyNames;
40+
41+
/**
42+
* Integration tests for observation instrumentation in {@link QianFanImageModel}.
43+
*
44+
* @author Geng Rong
45+
*/
46+
@SpringBootTest(classes = QianFanImageModelObservationIT.Config.class)
47+
@EnabledIfEnvironmentVariables(value = { @EnabledIfEnvironmentVariable(named = "QIANFAN_API_KEY", matches = ".+"),
48+
@EnabledIfEnvironmentVariable(named = "QIANFAN_SECRET_KEY", matches = ".+") })
49+
public class QianFanImageModelObservationIT {
50+
51+
@Autowired
52+
TestObservationRegistry observationRegistry;
53+
54+
@Autowired
55+
QianFanImageModel imageModel;
56+
57+
@Test
58+
void observationForImageOperation() {
59+
var options = QianFanImageOptions.builder()
60+
.withModel(QianFanImageApi.ImageModel.Stable_Diffusion_XL.getValue())
61+
.withHeight(1024)
62+
.withWidth(1024)
63+
.withStyle("Base")
64+
.build();
65+
66+
var instructions = "Here comes the sun";
67+
68+
ImagePrompt imagePrompt = new ImagePrompt(instructions, options);
69+
70+
ImageResponse imageResponse = imageModel.call(imagePrompt);
71+
assertThat(imageResponse.getResults()).hasSize(1);
72+
73+
TestObservationRegistryAssert.assertThat(observationRegistry)
74+
.doesNotHaveAnyRemainingCurrentObservation()
75+
.hasObservationWithNameEqualTo(DefaultImageModelObservationConvention.DEFAULT_NAME)
76+
.that()
77+
.hasContextualNameEqualTo("image " + QianFanImageApi.ImageModel.Stable_Diffusion_XL.getValue())
78+
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(),
79+
AiOperationType.IMAGE.value())
80+
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.AI_PROVIDER.asString(), AiProvider.QIANFAN.value())
81+
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.REQUEST_MODEL.asString(),
82+
QianFanImageApi.ImageModel.Stable_Diffusion_XL.getValue())
83+
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_IMAGE_SIZE.asString(), "1024x1024")
84+
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.REQUEST_IMAGE_STYLE.asString(), "Base")
85+
.hasBeenStarted()
86+
.hasBeenStopped();
87+
}
88+
89+
@SpringBootConfiguration
90+
static class Config {
91+
92+
@Bean
93+
public TestObservationRegistry observationRegistry() {
94+
return TestObservationRegistry.create();
95+
}
96+
97+
@Bean
98+
public QianFanImageApi qianFanImageApi() {
99+
return new QianFanImageApi(System.getenv("QIANFAN_API_KEY"), System.getenv("QIANFAN_SECRET_KEY"));
100+
}
101+
102+
@Bean
103+
public QianFanImageModel qianFanImageModel(QianFanImageApi qianFanImageApi,
104+
TestObservationRegistry observationRegistry) {
105+
return new QianFanImageModel(qianFanImageApi, QianFanImageOptions.builder().build(),
106+
RetryTemplate.defaultInstance(), observationRegistry);
107+
}
108+
109+
}
110+
111+
}

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/qianfan/QianFanAutoConfiguration.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
2020
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
2121
import org.springframework.ai.embedding.observation.EmbeddingModelObservationConvention;
22+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
2223
import org.springframework.ai.model.function.FunctionCallbackContext;
2324
import org.springframework.ai.qianfan.QianFanChatModel;
2425
import org.springframework.ai.qianfan.QianFanEmbeddingModel;
@@ -99,7 +100,8 @@ public QianFanEmbeddingModel qianFanEmbeddingModel(QianFanConnectionProperties c
99100
matchIfMissing = true)
100101
public QianFanImageModel qianFanImageModel(QianFanConnectionProperties commonProperties,
101102
QianFanImageProperties imageProperties, RestClient.Builder restClientBuilder, RetryTemplate retryTemplate,
102-
ResponseErrorHandler responseErrorHandler) {
103+
ResponseErrorHandler responseErrorHandler, ObjectProvider<ObservationRegistry> observationRegistry,
104+
ObjectProvider<ImageModelObservationConvention> observationConvention) {
103105

104106
String apiKey = StringUtils.hasText(imageProperties.getApiKey()) ? imageProperties.getApiKey()
105107
: commonProperties.getApiKey();
@@ -116,7 +118,12 @@ public QianFanImageModel qianFanImageModel(QianFanConnectionProperties commonPro
116118

117119
var qianFanImageApi = new QianFanImageApi(baseUrl, apiKey, secretKey, restClientBuilder, responseErrorHandler);
118120

119-
return new QianFanImageModel(qianFanImageApi, imageProperties.getOptions(), retryTemplate);
121+
var imageModel = new QianFanImageModel(qianFanImageApi, imageProperties.getOptions(), retryTemplate,
122+
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
123+
124+
observationConvention.ifAvailable(imageModel::setObservationConvention);
125+
126+
return imageModel;
120127
}
121128

122129
private QianFanApi qianFanApi(String baseUrl, String commonBaseUrl, String apiKey, String commonApiKey,

0 commit comments

Comments
 (0)