Fix Usage in Ollama ChatResponse

The Usage of Ollama ChatResponse was put into ChatGenerationMetadata as content filter metadata. The correct place should be in ChatResponseMetadata

Author: alexcheng1982

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java

@@ -18,6 +18,7 @@
 import java.util.Base64;
 import java.util.List;
 
+import org.springframework.ai.ollama.metadata.OllamaChatResponseMetadata;
 import reactor.core.publisher.Flux;
 
 import org.springframework.ai.chat.ChatClient;
@@ -27,7 +28,6 @@
 import org.springframework.ai.chat.messages.Message;
 import org.springframework.ai.chat.messages.MessageType;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
-import org.springframework.ai.chat.metadata.Usage;
 import org.springframework.ai.chat.prompt.ChatOptions;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.model.ModelOptionsUtils;
@@ -99,10 +99,9 @@ public ChatResponse call(Prompt prompt) {
 
 		var generator = new Generation(response.message().content());
 		if (response.promptEvalCount() != null && response.evalCount() != null) {
-			generator = generator
-				.withGenerationMetadata(ChatGenerationMetadata.from("unknown", extractUsage(response)));
+			generator = generator.withGenerationMetadata(ChatGenerationMetadata.from("unknown", null));
 		}
-		return new ChatResponse(List.of(generator));
+		return new ChatResponse(List.of(generator), OllamaChatResponseMetadata.from(response));
 	}
 
 	@Override
@@ -114,28 +113,12 @@ public Flux<ChatResponse> stream(Prompt prompt) {
 			Generation generation = (chunk.message() != null) ? new Generation(chunk.message().content())
 					: new Generation("");
 			if (Boolean.TRUE.equals(chunk.done())) {
-				generation = generation
-					.withGenerationMetadata(ChatGenerationMetadata.from("unknown", extractUsage(chunk)));
+				generation = generation.withGenerationMetadata(ChatGenerationMetadata.from("unknown", null));
 			}
-			return new ChatResponse(List.of(generation));
+			return new ChatResponse(List.of(generation), OllamaChatResponseMetadata.from(chunk));
 		});
 	}
 
-	private Usage extractUsage(OllamaApi.ChatResponse response) {
-		return new Usage() {
-
-			@Override
-			public Long getPromptTokens() {
-				return response.promptEvalCount().longValue();
-			}
-
-			@Override
-			public Long getGenerationTokens() {
-				return response.evalCount().longValue();
-			}
-		};
-	}
-
 	/**
 	 * Package access for testing.
 	 */

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaChatResponseMetadata.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2023 - 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.ollama.metadata;
+
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.ollama.api.OllamaApi;
+import org.springframework.util.Assert;
+
+/**
+ * {@link ChatResponseMetadata} implementation for {@literal Ollama}
+ *
+ * @see ChatResponseMetadata
+ * @author Fu Cheng
+ */
+public class OllamaChatResponseMetadata implements ChatResponseMetadata {
+
+	protected static final String AI_METADATA_STRING = "{ @type: %1$s, usage: %2$s, rateLimit: %3$s }";
+
+	public static OllamaChatResponseMetadata from(OllamaApi.ChatResponse response) {
+		Assert.notNull(response, "OllamaApi.ChatResponse must not be null");
+		Usage usage = OllamaUsage.from(response);
+		return new OllamaChatResponseMetadata(usage);
+	}
+
+	private final Usage usage;
+
+	protected OllamaChatResponseMetadata(Usage usage) {
+		this.usage = usage;
+	}
+
+	@Override
+	public Usage getUsage() {
+		return this.usage;
+	}
+
+	@Override
+	public String toString() {
+		return AI_METADATA_STRING.formatted(getClass().getTypeName(), getUsage(), getRateLimit());
+	}
+
+}

models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaUsage.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2023 - 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.ollama.metadata;
+
+import java.util.Optional;
+import org.springframework.ai.chat.metadata.Usage;
+import org.springframework.ai.ollama.api.OllamaApi;
+import org.springframework.util.Assert;
+
+/**
+ * {@link Usage} implementation for {@literal Ollama}
+ *
+ * @see Usage
+ * @author Fu Cheng
+ */
+public class OllamaUsage implements Usage {
+
+	protected static final String AI_USAGE_STRING = "{ promptTokens: %1$d, generationTokens: %2$d, totalTokens: %3$d }";
+
+	public static OllamaUsage from(OllamaApi.ChatResponse response) {
+		Assert.notNull(response, "OllamaApi.ChatResponse must not be null");
+		return new OllamaUsage(response);
+	}
+
+	private final OllamaApi.ChatResponse response;
+
+	public OllamaUsage(OllamaApi.ChatResponse response) {
+		this.response = response;
+	}
+
+	@Override
+	public Long getPromptTokens() {
+		return Optional.ofNullable(response.promptEvalCount()).map(Integer::longValue).orElse(0L);
+	}
+
+	@Override
+	public Long getGenerationTokens() {
+		return Optional.ofNullable(response.evalCount()).map(Integer::longValue).orElse(0L);
+	}
+
+	@Override
+	public String toString() {
+		return AI_USAGE_STRING.formatted(getPromptTokens(), getGenerationTokens(), getTotalTokens());
+	}
+
+}

models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java

@@ -26,6 +26,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.metadata.Usage;
 import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
 import org.springframework.ai.chat.messages.AssistantMessage;
 import org.testcontainers.containers.GenericContainer;
@@ -105,6 +106,18 @@ void roleTest() {
 
 	}
 
+	@Test
+	void usageTest() {
+		Prompt prompt = new Prompt("Tell me a joke");
+		ChatResponse response = client.call(prompt);
+		Usage usage = response.getMetadata().getUsage();
+
+		assertThat(usage).isNotNull();
+		assertThat(usage.getPromptTokens()).isPositive();
+		assertThat(usage.getGenerationTokens()).isPositive();
+		assertThat(usage.getTotalTokens()).isPositive();
+	}
+
 	@Test
 	void outputParser() {
 		DefaultConversionService conversionService = new DefaultConversionService();