Skip to content

Commit f282131

Browse files
quaffquaff
committed
Instrument observation and retry for ImageModel implementations
Signed-off-by: Yanming Zhou <[email protected]>
1 parent c2dfedb commit f282131

File tree

13 files changed

+267
-72
lines changed

13 files changed

+267
-72
lines changed

‎auto-configurations/models/spring-ai-autoconfigure-model-azure-openai/src/main/java/org/springframework/ai/model/azure/openai/autoconfigure/AzureOpenAiImageAutoConfiguration.java‎

Lines changed: 13 additions & 3 deletions
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-2026 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,16 +18,20 @@
1818

1919
import com.azure.ai.openai.OpenAIClientBuilder;
2020

21+
import io.micrometer.observation.ObservationRegistry;
2122
import org.springframework.ai.azure.openai.AzureOpenAiImageModel;
23+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
2224
import org.springframework.ai.model.SpringAIModelProperties;
2325
import org.springframework.ai.model.SpringAIModels;
26+
import org.springframework.beans.factory.ObjectProvider;
2427
import org.springframework.boot.autoconfigure.AutoConfiguration;
2528
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2629
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
2730
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
2831
import org.springframework.boot.context.properties.EnableConfigurationProperties;
2932
import org.springframework.context.annotation.Bean;
3033
import org.springframework.context.annotation.Import;
34+
import org.springframework.core.retry.RetryTemplate;
3135

3236
/**
3337
* {@link AutoConfiguration Auto-configuration} for Azure OpenAI.
@@ -36,6 +40,7 @@
3640
* @author Soby Chacko
3741
* @author Manuel Andreo Garcia
3842
* @author Ilayaperumal Gopinathan
43+
* @author Yanming Zhou
3944
*/
4045
@AutoConfiguration
4146
@ConditionalOnClass(AzureOpenAiImageModel.class)
@@ -48,9 +53,14 @@ public class AzureOpenAiImageAutoConfiguration {
4853
@Bean
4954
@ConditionalOnMissingBean
5055
public AzureOpenAiImageModel azureOpenAiImageModel(OpenAIClientBuilder openAIClientBuilder,
51-
AzureOpenAiImageOptionsProperties imageProperties) {
56+
AzureOpenAiImageOptionsProperties imageProperties, RetryTemplate retryTemplate,
57+
ObjectProvider<ObservationRegistry> observationRegistry,
58+
ObjectProvider<ImageModelObservationConvention> observationConvention) {
5259

53-
return new AzureOpenAiImageModel(openAIClientBuilder.buildClient(), imageProperties.getOptions());
60+
var imageModel = new AzureOpenAiImageModel(openAIClientBuilder.buildClient(), imageProperties.getOptions(),
61+
retryTemplate, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
62+
observationConvention.ifAvailable(imageModel::setObservationConvention);
63+
return imageModel;
5464
}
5565

5666
}

‎auto-configurations/models/spring-ai-autoconfigure-model-openai-sdk/src/main/java/org/springframework/ai/model/openaisdk/autoconfigure/OpenAiSdkImageAutoConfiguration.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2025 the original author or authors.
2+
* Copyright 2023-2026 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.
@@ -30,11 +30,13 @@
3030
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3131
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3232
import org.springframework.context.annotation.Bean;
33+
import org.springframework.core.retry.RetryTemplate;
3334

3435
/**
3536
* Image {@link AutoConfiguration Auto-configuration} for OpenAI.
3637
*
3738
* @author Christian Tzolov
39+
* @author Yanming Zhou
3840
*/
3941
@AutoConfiguration
4042
@ConditionalOnProperty(name = SpringAIModelProperties.IMAGE_MODEL, havingValue = SpringAIModels.OPENAI_SDK,
@@ -45,11 +47,13 @@ public class OpenAiSdkImageAutoConfiguration {
4547
@Bean
4648
@ConditionalOnMissingBean
4749
public OpenAiSdkImageModel openAiImageModel(OpenAiSdkConnectionProperties commonProperties,
48-
OpenAiSdkImageProperties imageProperties, ObjectProvider<ObservationRegistry> observationRegistry,
50+
OpenAiSdkImageProperties imageProperties, RetryTemplate retryTemplate,
51+
ObjectProvider<ObservationRegistry> observationRegistry,
4952
ObjectProvider<ImageModelObservationConvention> observationConvention) {
5053

5154
var imageModel = new OpenAiSdkImageModel(openAiClient(commonProperties, imageProperties),
52-
imageProperties.getOptions(), observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
55+
imageProperties.getOptions(), retryTemplate,
56+
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
5357

5458
observationConvention.ifAvailable(imageModel::setObservationConvention);
5559

‎auto-configurations/models/spring-ai-autoconfigure-model-stability-ai/src/main/java/org/springframework/ai/model/stabilityai/autoconfigure/StabilityAiImageAutoConfiguration.java‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2025 the original author or authors.
2+
* Copyright 2023-2026 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.
@@ -16,6 +16,8 @@
1616

1717
package org.springframework.ai.model.stabilityai.autoconfigure;
1818

19+
import io.micrometer.observation.ObservationRegistry;
20+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
1921
import org.springframework.ai.model.SpringAIModelProperties;
2022
import org.springframework.ai.model.SpringAIModels;
2123
import org.springframework.ai.stabilityai.StabilityAiImageModel;
@@ -28,6 +30,7 @@
2830
import org.springframework.boot.context.properties.EnableConfigurationProperties;
2931
import org.springframework.boot.restclient.autoconfigure.RestClientAutoConfiguration;
3032
import org.springframework.context.annotation.Bean;
33+
import org.springframework.core.retry.RetryTemplate;
3134
import org.springframework.util.Assert;
3235
import org.springframework.util.StringUtils;
3336
import org.springframework.web.client.RestClient;
@@ -38,6 +41,7 @@
3841
* @author Mark Pollack
3942
* @author Christian Tzolov
4043
* @author Ilayaperumal Gopinathan
44+
* @author Yanming Zhou
4145
* @since 0.8.0
4246
*/
4347
@AutoConfiguration(after = RestClientAutoConfiguration.class)
@@ -68,8 +72,13 @@ public StabilityAiApi stabilityAiApi(StabilityAiConnectionProperties commonPrope
6872
@Bean
6973
@ConditionalOnMissingBean
7074
public StabilityAiImageModel stabilityAiImageModel(StabilityAiApi stabilityAiApi,
71-
StabilityAiImageProperties stabilityAiImageProperties) {
72-
return new StabilityAiImageModel(stabilityAiApi, stabilityAiImageProperties.getOptions());
75+
StabilityAiImageProperties stabilityAiImageProperties, RetryTemplate retryTemplate,
76+
ObjectProvider<ObservationRegistry> observationRegistry,
77+
ObjectProvider<ImageModelObservationConvention> observationConvention) {
78+
var imageModel = new StabilityAiImageModel(stabilityAiApi, stabilityAiImageProperties.getOptions(),
79+
retryTemplate, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
80+
observationConvention.ifAvailable(imageModel::setObservationConvention);
81+
return imageModel;
7382
}
7483

7584
}

‎auto-configurations/models/spring-ai-autoconfigure-model-zhipuai/src/main/java/org/springframework/ai/model/zhipuai/autoconfigure/ZhiPuAiImageAutoConfiguration.java‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2023-2025 the original author or authors.
2+
* Copyright 2023-2026 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.
@@ -16,6 +16,8 @@
1616

1717
package org.springframework.ai.model.zhipuai.autoconfigure;
1818

19+
import io.micrometer.observation.ObservationRegistry;
20+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
1921
import org.springframework.ai.model.SpringAIModelProperties;
2022
import org.springframework.ai.model.SpringAIModels;
2123
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
@@ -41,6 +43,7 @@
4143
*
4244
* @author Geng Rong
4345
* @author Ilayaperumal Gopinathan
46+
* @author Yanming Zhou
4447
*/
4548
@AutoConfiguration(after = { RestClientAutoConfiguration.class, SpringAiRetryAutoConfiguration.class })
4649
@ConditionalOnClass(ZhiPuAiApi.class)
@@ -53,7 +56,9 @@ public class ZhiPuAiImageAutoConfiguration {
5356
@ConditionalOnMissingBean
5457
public ZhiPuAiImageModel zhiPuAiImageModel(ZhiPuAiConnectionProperties commonProperties,
5558
ZhiPuAiImageProperties imageProperties, ObjectProvider<RestClient.Builder> restClientBuilderProvider,
56-
RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) {
59+
RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler,
60+
ObjectProvider<ObservationRegistry> observationRegistry,
61+
ObjectProvider<ImageModelObservationConvention> observationConvention) {
5762

5863
String apiKey = StringUtils.hasText(imageProperties.getApiKey()) ? imageProperties.getApiKey()
5964
: commonProperties.getApiKey();
@@ -64,11 +69,13 @@ public ZhiPuAiImageModel zhiPuAiImageModel(ZhiPuAiConnectionProperties commonPro
6469
Assert.hasText(apiKey, "ZhiPuAI API key must be set");
6570
Assert.hasText(baseUrl, "ZhiPuAI base URL must be set");
6671

67-
// TODO add ZhiPuAiApi support for image
6872
var zhiPuAiImageApi = new ZhiPuAiImageApi(baseUrl, apiKey,
6973
restClientBuilderProvider.getIfAvailable(RestClient::builder), responseErrorHandler);
7074

71-
return new ZhiPuAiImageModel(zhiPuAiImageApi, imageProperties.getOptions(), retryTemplate);
75+
var imageModel = new ZhiPuAiImageModel(zhiPuAiImageApi, imageProperties.getOptions(), retryTemplate,
76+
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP));
77+
observationConvention.ifAvailable(imageModel::setObservationConvention);
78+
return imageModel;
7279
}
7380

7481
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@
4747
<version>${project.parent.version}</version>
4848
</dependency>
4949

50+
<dependency>
51+
<groupId>org.springframework.ai</groupId>
52+
<artifactId>spring-ai-retry</artifactId>
53+
<version>${project.parent.version}</version>
54+
</dependency>
55+
5056
<dependency>
5157
<groupId>com.azure</groupId>
5258
<artifactId>azure-ai-openai</artifactId>

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

Lines changed: 61 additions & 12 deletions
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-2026 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.
@@ -29,6 +29,7 @@
2929
import com.fasterxml.jackson.databind.ObjectMapper;
3030
import com.fasterxml.jackson.databind.SerializationFeature;
3131
import com.fasterxml.jackson.databind.json.JsonMapper;
32+
import io.micrometer.observation.ObservationRegistry;
3233
import org.slf4j.Logger;
3334
import org.slf4j.LoggerFactory;
3435

@@ -40,8 +41,14 @@
4041
import org.springframework.ai.image.ImagePrompt;
4142
import org.springframework.ai.image.ImageResponse;
4243
import org.springframework.ai.image.ImageResponseMetadata;
44+
import org.springframework.ai.image.observation.ImageModelObservationContext;
45+
import org.springframework.ai.image.observation.ImageModelObservationConvention;
46+
import org.springframework.ai.image.observation.ImageModelObservationDocumentation;
4347
import org.springframework.ai.model.ModelOptionsUtils;
48+
import org.springframework.ai.observation.conventions.AiProvider;
49+
import org.springframework.ai.retry.RetryUtils;
4450
import org.springframework.ai.util.JacksonUtils;
51+
import org.springframework.core.retry.RetryTemplate;
4552
import org.springframework.util.Assert;
4653

4754
/**
@@ -50,6 +57,7 @@
5057
*
5158
* @author Benoit Moussaud
5259
* @author Sebastien Deleuze
60+
* @author Yanming Zhou
5361
* @see ImageModel
5462
* @see com.azure.ai.openai.OpenAIClient
5563
* @since 1.0.0
@@ -66,20 +74,49 @@ public class AzureOpenAiImageModel implements ImageModel {
6674

6775
private final ObjectMapper objectMapper;
6876

77+
private final RetryTemplate retryTemplate;
78+
79+
private final ObservationRegistry observationRegistry;
80+
81+
private ImageModelObservationConvention observationConvention = DEFAULT_OBSERVATION_CONVENTION;
82+
6983
public AzureOpenAiImageModel(OpenAIClient openAIClient) {
7084
this(openAIClient, AzureOpenAiImageOptions.builder().deploymentName(DEFAULT_DEPLOYMENT_NAME).build());
7185
}
7286

7387
public AzureOpenAiImageModel(OpenAIClient microsoftOpenAiClient, AzureOpenAiImageOptions options) {
88+
this(microsoftOpenAiClient, options, RetryUtils.DEFAULT_RETRY_TEMPLATE, ObservationRegistry.NOOP);
89+
}
90+
91+
public AzureOpenAiImageModel(OpenAIClient microsoftOpenAiClient, AzureOpenAiImageOptions options,
92+
RetryTemplate retryTemplate) {
93+
this(microsoftOpenAiClient, options, retryTemplate, ObservationRegistry.NOOP);
94+
}
95+
96+
public AzureOpenAiImageModel(OpenAIClient microsoftOpenAiClient, AzureOpenAiImageOptions options,
97+
RetryTemplate retryTemplate, ObservationRegistry observationRegistry) {
7498
Assert.notNull(microsoftOpenAiClient, "com.azure.ai.openai.OpenAIClient must not be null");
7599
Assert.notNull(options, "AzureOpenAiChatOptions must not be null");
100+
Assert.notNull(retryTemplate, "retryTemplate must not be null");
101+
Assert.notNull(observationRegistry, "observationRegistry must not be null");
76102
this.openAIClient = microsoftOpenAiClient;
77103
this.defaultOptions = options;
78104
this.objectMapper = JsonMapper.builder()
79105
.addModules(JacksonUtils.instantiateAvailableModules())
80106
.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
81107
.disable(SerializationFeature.FAIL_ON_EMPTY_BEANS)
82108
.build();
109+
this.retryTemplate = retryTemplate;
110+
this.observationRegistry = observationRegistry;
111+
}
112+
113+
/**
114+
* Use the provided convention for reporting observation data
115+
* @param observationConvention The provided convention
116+
*/
117+
public void setObservationConvention(ImageModelObservationConvention observationConvention) {
118+
Assert.notNull(observationConvention, "observationConvention cannot be null");
119+
this.observationConvention = observationConvention;
83120
}
84121

85122
public AzureOpenAiImageOptions getDefaultOptions() {
@@ -95,20 +132,32 @@ public ImageResponse call(ImagePrompt imagePrompt) {
95132
toPrettyJson(imageGenerationOptions));
96133
}
97134

98-
var images = this.openAIClient.getImageGenerations(deploymentOrModelName, imageGenerationOptions);
135+
var observationContext = ImageModelObservationContext.builder()
136+
.imagePrompt(imagePrompt)
137+
.provider(AiProvider.AZURE_OPENAI.name())
138+
.build();
99139

100-
if (logger.isTraceEnabled()) {
101-
logger.trace("Azure ImageGenerations: {}", toPrettyJson(images));
102-
}
140+
return ImageModelObservationDocumentation.IMAGE_MODEL_OPERATION
141+
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
142+
this.observationRegistry)
143+
.observe(() -> {
144+
var images = RetryUtils.execute(this.retryTemplate,
145+
() -> this.openAIClient.getImageGenerations(deploymentOrModelName, imageGenerationOptions));
146+
147+
if (logger.isTraceEnabled()) {
148+
logger.trace("Azure ImageGenerations: {}", toPrettyJson(images));
149+
}
150+
151+
List<ImageGeneration> imageGenerations = images.getData().stream().map(entry -> {
152+
var image = new Image(entry.getUrl(), entry.getBase64Data());
153+
var metadata = new AzureOpenAiImageGenerationMetadata(entry.getRevisedPrompt());
154+
return new ImageGeneration(image, metadata);
155+
}).toList();
103156

104-
List<ImageGeneration> imageGenerations = images.getData().stream().map(entry -> {
105-
var image = new Image(entry.getUrl(), entry.getBase64Data());
106-
var metadata = new AzureOpenAiImageGenerationMetadata(entry.getRevisedPrompt());
107-
return new ImageGeneration(image, metadata);
108-
}).toList();
157+
ImageResponseMetadata openAiImageResponseMetadata = AzureOpenAiImageResponseMetadata.from(images);
158+
return new ImageResponse(imageGenerations, openAiImageResponseMetadata);
159+
});
109160

110-
ImageResponseMetadata openAiImageResponseMetadata = AzureOpenAiImageResponseMetadata.from(images);
111-
return new ImageResponse(imageGenerations, openAiImageResponseMetadata);
112161
}
113162

114163
private String toPrettyJson(Object object) {

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@
4747
<version>${project.parent.version}</version>
4848
</dependency>
4949

50+
<dependency>
51+
<groupId>org.springframework.ai</groupId>
52+
<artifactId>spring-ai-retry</artifactId>
53+
<version>${project.parent.version}</version>
54+
</dependency>
55+
5056
<dependency>
5157
<groupId>com.openai</groupId>
5258
<artifactId>openai-java</artifactId>

0 commit comments

Comments
 (0)