From 2129547923b053ef7258695f4b4f4aa074208b2d Mon Sep 17 00:00:00 2001 From: Fu Cheng Date: Thu, 28 Mar 2024 15:11:20 +0800 Subject: [PATCH] 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 --- .../ai/ollama/OllamaChatClient.java | 27 ++------- .../metadata/OllamaChatResponseMetadata.java | 55 +++++++++++++++++ .../ai/ollama/metadata/OllamaUsage.java | 59 +++++++++++++++++++ .../ai/ollama/OllamaChatClientIT.java | 13 ++++ 4 files changed, 132 insertions(+), 22 deletions(-) create mode 100644 models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaChatResponseMetadata.java create mode 100644 models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaUsage.java diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java index 4efbc3a77..7629bc0fb 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatClient.java @@ -18,6 +18,7 @@ package org.springframework.ai.ollama; 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.StreamingChatClient; 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 class OllamaChatClient implements ChatClient, StreamingChatClient { 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 class OllamaChatClient implements ChatClient, StreamingChatClient { 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. */ diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaChatResponseMetadata.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaChatResponseMetadata.java new file mode 100644 index 000000000..906cf1075 --- /dev/null +++ b/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()); + } + +} diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaUsage.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaUsage.java new file mode 100644 index 000000000..a437557d5 --- /dev/null +++ b/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()); + } + +} diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java index b76cae1dd..4ee531add 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatClientIT.java @@ -26,6 +26,7 @@ import org.apache.commons.logging.LogFactory; 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 @@ class OllamaChatClientIT { } + @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();