Skip to content

Commit d4e36db

Browse files
committed
Address review comments
1 parent 8be1e42 commit d4e36db

File tree

1 file changed

+32
-16
lines changed

1 file changed

+32
-16
lines changed

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

Lines changed: 32 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,22 @@
8181
*/
8282
public class OllamaChatModel extends AbstractToolCallSupport implements ChatModel {
8383

84+
private static final String DONE = "done";
85+
86+
private static final String METADATA_PROMPT_EVAL_COUNT = "prompt-eval-count";
87+
88+
private static final String METADATA_EVAL_COUNT = "eval-count";
89+
90+
private static final String METADATA_CREATED_AT = "created-at";
91+
92+
private static final String METADATA_TOTAL_DURATION = "total-duration";
93+
94+
private static final String METADATA_LOAD_DURATION = "load-duration";
95+
96+
private static final String METADATA_PROMPT_EVAL_DURATION = "prompt-eval-duration";
97+
98+
private static final String METADATA_EVAL_DURATION = "eval-duration";
99+
84100
private static final ChatModelObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultChatModelObservationConvention();
85101

86102
private final OllamaApi chatApi;
@@ -126,18 +142,18 @@ static ChatResponseMetadata from(OllamaApi.ChatResponse response, ChatResponse p
126142
Duration totalDuration = response.getTotalDuration();
127143

128144
if (previousChatResponse != null && previousChatResponse.getMetadata() != null) {
129-
if (previousChatResponse.getMetadata().get("eval-duration") != null) {
130-
evalDuration = evalDuration.plus(previousChatResponse.getMetadata().get("eval-duration"));
145+
if (previousChatResponse.getMetadata().get(METADATA_EVAL_DURATION) != null) {
146+
evalDuration = evalDuration.plus(previousChatResponse.getMetadata().get(METADATA_EVAL_DURATION));
131147
}
132-
if (previousChatResponse.getMetadata().get("prompt-eval-duration") != null) {
148+
if (previousChatResponse.getMetadata().get(METADATA_PROMPT_EVAL_DURATION) != null) {
133149
promptEvalDuration = promptEvalDuration
134-
.plus(previousChatResponse.getMetadata().get("prompt-eval-duration"));
150+
.plus(previousChatResponse.getMetadata().get(METADATA_PROMPT_EVAL_DURATION));
135151
}
136-
if (previousChatResponse.getMetadata().get("load-duration") != null) {
137-
loadDuration = loadDuration.plus(previousChatResponse.getMetadata().get("load-duration"));
152+
if (previousChatResponse.getMetadata().get(METADATA_LOAD_DURATION) != null) {
153+
loadDuration = loadDuration.plus(previousChatResponse.getMetadata().get(METADATA_LOAD_DURATION));
138154
}
139-
if (previousChatResponse.getMetadata().get("total-duration") != null) {
140-
totalDuration = totalDuration.plus(previousChatResponse.getMetadata().get("total-duration"));
155+
if (previousChatResponse.getMetadata().get(METADATA_TOTAL_DURATION) != null) {
156+
totalDuration = totalDuration.plus(previousChatResponse.getMetadata().get(METADATA_TOTAL_DURATION));
141157
}
142158
if (previousChatResponse.getMetadata().getUsage() != null) {
143159
promptTokens += previousChatResponse.getMetadata().getUsage().getPromptTokens();
@@ -151,14 +167,14 @@ static ChatResponseMetadata from(OllamaApi.ChatResponse response, ChatResponse p
151167
return ChatResponseMetadata.builder()
152168
.withUsage(aggregatedUsage)
153169
.withModel(response.model())
154-
.withKeyValue("created-at", response.createdAt())
155-
.withKeyValue("eval-duration", evalDuration)
156-
.withKeyValue("eval-count", aggregatedUsage.getGenerationTokens().intValue())
157-
.withKeyValue("load-duration", loadDuration)
158-
.withKeyValue("prompt-eval-duration", promptEvalDuration)
159-
.withKeyValue("prompt-eval-count", aggregatedUsage.getPromptTokens().intValue())
160-
.withKeyValue("total-duration", totalDuration)
161-
.withKeyValue("done", response.done())
170+
.withKeyValue(METADATA_CREATED_AT, response.createdAt())
171+
.withKeyValue(METADATA_EVAL_DURATION, evalDuration)
172+
.withKeyValue(METADATA_EVAL_COUNT, aggregatedUsage.getGenerationTokens().intValue())
173+
.withKeyValue(METADATA_LOAD_DURATION, loadDuration)
174+
.withKeyValue(METADATA_PROMPT_EVAL_DURATION, promptEvalDuration)
175+
.withKeyValue(METADATA_PROMPT_EVAL_COUNT, aggregatedUsage.getPromptTokens().intValue())
176+
.withKeyValue(METADATA_TOTAL_DURATION, totalDuration)
177+
.withKeyValue(DONE, response.done())
162178
.build();
163179
}
164180

0 commit comments

Comments
 (0)