Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix Usage in Ollama ChatResponse

Browse Source 
The Usage of Ollama ChatResponse was put into ChatGenerationMetadata as content filter metadata. The correct place should be in ChatResponseMetadata
This commit is contained in:
Fu Cheng
2024-03-28 15:11:20 +08:00
committed by Christian Tzolov
parent 9236913696
commit 2129547923
4 changed files with 132 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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.
*/

55
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaChatResponseMetadata.java Normal file
View File

@@ -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());
}
}

59
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/metadata/OllamaUsage.java Normal file
View File

@@ -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());
}
}

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

@@ -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();