fix: missing micrometer metrics for images

auto-configurations/models/image/observation/spring-ai-autoconfigure-model-image-observation/pom.xml

@@ -66,6 +66,18 @@
 			<artifactId>mockito-core</artifactId>
 			<scope>test</scope>
 		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-micrometer-metrics</artifactId>
+			<scope>test</scope>
+		</dependency>
+
+		<dependency>
+			<groupId>org.springframework.boot</groupId>
+			<artifactId>spring-boot-micrometer-observation</artifactId>
+			<scope>test</scope>
+		</dependency>
 	</dependencies>
 
 </project>

auto-configurations/models/image/observation/spring-ai-autoconfigure-model-image-observation/src/main/java/org/springframework/ai/model/image/observation/autoconfigure/ImageObservationAutoConfiguration.java

@@ -16,14 +16,17 @@
 
 package org.springframework.ai.model.image.observation.autoconfigure;
 
+import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.tracing.Tracer;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
 import org.springframework.ai.image.ImageModel;
+import org.springframework.ai.image.observation.ImageModelMeterObservationHandler;
 import org.springframework.ai.image.observation.ImageModelObservationContext;
 import org.springframework.ai.image.observation.ImageModelPromptContentObservationHandler;
 import org.springframework.ai.observation.TracingAwareLoggingObservationHandler;
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -41,7 +44,8 @@
  * @author Jonatan Ivanov
  * @since 1.0.0
  */
-@AutoConfiguration
+@AutoConfiguration(
+		afterName = "org.springframework.boot.micrometer.metrics.autoconfigure.CompositeMeterRegistryAutoConfiguration")
 @ConditionalOnClass(ImageModel.class)
 @EnableConfigurationProperties(ImageObservationProperties.class)
 public class ImageObservationAutoConfiguration {
@@ -53,6 +57,13 @@ private static void logPromptContentWarning() {
 				"You have enabled logging out the image prompt content with the risk of exposing sensitive or private information. Please, be careful!");
 	}
 
+	@Bean
+	@ConditionalOnMissingBean
+	@ConditionalOnBean(MeterRegistry.class)
+	ImageModelMeterObservationHandler imageModelMeterObservationHandler(ObjectProvider<MeterRegistry> meterRegistry) {
+		return new ImageModelMeterObservationHandler(meterRegistry.getObject());
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	@ConditionalOnClass(Tracer.class)
 	@ConditionalOnBean(Tracer.class)

auto-configurations/models/image/observation/spring-ai-autoconfigure-model-image-observation/src/test/java/org/springframework/ai/model/image/observation/autoconfigure/ImageObservationAutoConfigurationTests.java

@@ -16,10 +16,12 @@
 
 package org.springframework.ai.model.image.observation.autoconfigure;
 
+import io.micrometer.core.instrument.composite.CompositeMeterRegistry;
 import io.micrometer.tracing.Tracer;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 
+import org.springframework.ai.image.observation.ImageModelMeterObservationHandler;
 import org.springframework.ai.image.observation.ImageModelObservationContext;
 import org.springframework.ai.image.observation.ImageModelPromptContentObservationHandler;
 import org.springframework.ai.observation.TracingAwareLoggingObservationHandler;
@@ -46,6 +48,17 @@ class ImageObservationAutoConfigurationTests {
 	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
 		.withConfiguration(AutoConfigurations.of(ImageObservationAutoConfiguration.class));
 
+	@Test
+	void meterObservationHandlerEnabled() {
+		this.contextRunner.withBean(CompositeMeterRegistry.class)
+			.run(context -> assertThat(context).hasSingleBean(ImageModelMeterObservationHandler.class));
+	}
+
+	@Test
+	void meterObservationHandlerDisabled() {
+		this.contextRunner.run(context -> assertThat(context).doesNotHaveBean(ImageModelMeterObservationHandler.class));
+	}
+
 	@Test
 	void imageModelPromptContentHandlerNoTracer() {
 		this.contextRunner.withClassLoader(new FilteredClassLoader(Tracer.class))

spring-ai-model/src/main/java/org/springframework/ai/image/observation/ImageModelMeterObservationHandler.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2023-present 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.image.observation;
+
+import java.util.Objects;
+
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationHandler;
+
+import org.springframework.ai.model.observation.ModelUsageMetricsGenerator;
+
+/**
+ * Handler for generating metrics from image model observations.
+ *
+ * @author Olivier Le Quellec
+ * @since 1.0.0
+ */
+public class ImageModelMeterObservationHandler implements ObservationHandler<ImageModelObservationContext> {
+
+	private final MeterRegistry meterRegistry;
+
+	public ImageModelMeterObservationHandler(MeterRegistry meterRegistry) {
+		this.meterRegistry = meterRegistry;
+	}
+
+	@Override
+	public void onStop(ImageModelObservationContext context) {
+		if (Objects.nonNull(context.getResponse()) && Objects.nonNull(context.getResponse().getMetadata())
+				&& Objects.nonNull(context.getResponse().getMetadata().getUsage())) {
+			ModelUsageMetricsGenerator.generate(context.getResponse().getMetadata().getUsage(), context,
+					this.meterRegistry);
+		}
+	}
+
+	@Override
+	public boolean supportsContext(Observation.Context context) {
+		return context instanceof ImageModelObservationContext;
+	}
+
+}

spring-ai-model/src/main/java/org/springframework/ai/model/observation/ModelUsageMetricsGenerator.java

@@ -18,6 +18,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 import io.micrometer.common.KeyValue;
 import io.micrometer.core.instrument.Counter;
@@ -45,7 +46,7 @@ private ModelUsageMetricsGenerator() {
 
 	public static void generate(Usage usage, Observation.Context context, MeterRegistry meterRegistry) {
 
-		if (usage.getPromptTokens() != null) {
+		if (Objects.nonNull(usage.getPromptTokens())) {
 			Counter.builder(AiObservationMetricNames.TOKEN_USAGE.value())
 				.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.INPUT.value())
 				.description(DESCRIPTION)
@@ -54,7 +55,7 @@ public static void generate(Usage usage, Observation.Context context, MeterRegis
 				.increment(usage.getPromptTokens());
 		}
 
-		if (usage.getCompletionTokens() != null) {
+		if (Objects.nonNull(usage.getCompletionTokens())) {
 			Counter.builder(AiObservationMetricNames.TOKEN_USAGE.value())
 				.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.OUTPUT.value())
 				.description(DESCRIPTION)
@@ -63,7 +64,7 @@ public static void generate(Usage usage, Observation.Context context, MeterRegis
 				.increment(usage.getCompletionTokens());
 		}
 
-		if (usage.getTotalTokens() != null) {
+		if (Objects.nonNull(usage.getTotalTokens())) {
 			Counter.builder(AiObservationMetricNames.TOKEN_USAGE.value())
 				.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.TOTAL.value())
 				.description(DESCRIPTION)

spring-ai-model/src/test/java/org/springframework/ai/image/observation/ImageModelMeterObservationHandlerTests.java

@@ -0,0 +1,146 @@
+/*
+ * Copyright 2023-present 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.image.observation;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.image.ImageMessage;
+import org.springframework.ai.image.ImageOptions;
+import org.springframework.ai.image.ImageOptionsBuilder;
+import org.springframework.ai.image.ImagePrompt;
+import org.springframework.ai.image.ImageResponse;
+import org.springframework.ai.image.ImageResponseMetadata;
+import org.springframework.ai.observation.conventions.AiObservationMetricAttributes;
+import org.springframework.ai.observation.conventions.AiObservationMetricNames;
+import org.springframework.ai.observation.conventions.AiOperationType;
+import org.springframework.ai.observation.conventions.AiTokenType;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.springframework.ai.image.observation.ImageModelObservationDocumentation.LowCardinalityKeyNames;
+
+/**
+ * Unit tests for {@link ImageModelMeterObservationHandler}.
+ *
+ * @author Olivier Le Quellec
+ */
+class ImageModelMeterObservationHandlerTests {
+
+	private MeterRegistry meterRegistry;
+
+	private ObservationRegistry observationRegistry;
+
+	@BeforeEach
+	void setUp() {
+		this.meterRegistry = new SimpleMeterRegistry();
+		this.observationRegistry = ObservationRegistry.create();
+		this.observationRegistry.observationConfig()
+			.observationHandler(new ImageModelMeterObservationHandler(this.meterRegistry));
+	}
+
+	@Test
+	void shouldCreateAllMetersDuringAnObservation() {
+		var observationContext = generateObservationContext();
+		var observation = Observation
+			.createNotStarted(new DefaultImageModelObservationConvention(), () -> observationContext,
+					this.observationRegistry)
+			.start();
+
+		observationContext.setResponse(
+				new ImageResponse(List.of(), new ImageResponseMetadata(System.currentTimeMillis(), new TestUsage())));
+
+		observation.stop();
+
+		assertThat(this.meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value()).meters()).hasSize(3);
+		assertThat(this.meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
+			.tag(LowCardinalityKeyNames.AI_OPERATION_TYPE.asString(), AiOperationType.IMAGE.value())
+			.tag(LowCardinalityKeyNames.AI_PROVIDER.asString(), "superprovider")
+			.tag(LowCardinalityKeyNames.REQUEST_MODEL.asString(), "dall-e-3")
+			.meters()).hasSize(3);
+		assertThat(this.meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
+			.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.INPUT.value())
+			.meters()).hasSize(1);
+		assertThat(this.meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
+			.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.OUTPUT.value())
+			.meters()).hasSize(1);
+		assertThat(this.meterRegistry.get(AiObservationMetricNames.TOKEN_USAGE.value())
+			.tag(AiObservationMetricAttributes.TOKEN_TYPE.value(), AiTokenType.TOTAL.value())
+			.meters()).hasSize(1);
+	}
+
+	@Test
+	void shouldNotCreateMetersWhenResponseIsNull() {
+		var observationContext = generateObservationContext();
+		var observation = Observation
+			.createNotStarted(new DefaultImageModelObservationConvention(), () -> observationContext,
+					this.observationRegistry)
+			.start();
+
+		observation.stop();
+
+		assertThat(this.meterRegistry.find(AiObservationMetricNames.TOKEN_USAGE.value()).meters()).isEmpty();
+	}
+
+	private ImageModelObservationContext generateObservationContext() {
+		return ImageModelObservationContext.builder()
+			.imagePrompt(generateImagePrompt(ImageOptionsBuilder.builder().model("dall-e-3").build()))
+			.provider("superprovider")
+			.build();
+	}
+
+	private ImagePrompt generateImagePrompt(ImageOptions imageOptions) {
+		return new ImagePrompt(new ImageMessage("a cute cat"), imageOptions);
+	}
+
+	static class TestUsage implements Usage {
+
+		@Override
+		public Integer getPromptTokens() {
+			return 1000;
+		}
+
+		@Override
+		public Integer getCompletionTokens() {
+			return 0;
+		}
+
+		@Override
+		public Integer getTotalTokens() {
+			return 1000;
+		}
+
+		@Override
+		public Map<String, Integer> getNativeUsage() {
+			Map<String, Integer> usage = new HashMap<>();
+			usage.put("promptTokens", getPromptTokens());
+			usage.put("completionTokens", getCompletionTokens());
+			usage.put("totalTokens", getTotalTokens());
+			return usage;
+		}
+
+	}
+
+}