From a555d00f747fe0abc2d2eb357890c79ca44f4f69 Mon Sep 17 00:00:00 2001 From: extio-de Date: Fri, 27 Jun 2025 21:19:26 +0200 Subject: [PATCH] Get token usage from OpenAI Api if available --- .../client/oai/textcompletion/Choice.java | 2 +- .../textcompletion/CompletionResponse.java | 17 ++++++- .../textcompletion/TextCompletionClient.java | 35 ++++++++++----- .../client/oai/textcompletion/Usage.java | 45 +++++++++++++++++++ 4 files changed, 85 insertions(+), 14 deletions(-) create mode 100644 src/main/java/de/extio/lmlib/client/oai/textcompletion/Usage.java diff --git a/src/main/java/de/extio/lmlib/client/oai/textcompletion/Choice.java b/src/main/java/de/extio/lmlib/client/oai/textcompletion/Choice.java index 7794eec..489d612 100755 --- a/src/main/java/de/extio/lmlib/client/oai/textcompletion/Choice.java +++ b/src/main/java/de/extio/lmlib/client/oai/textcompletion/Choice.java @@ -15,7 +15,7 @@ public final class Choice { private Double score; private String finishReason; - + @JsonProperty("text") public String getText() { return this.text; diff --git a/src/main/java/de/extio/lmlib/client/oai/textcompletion/CompletionResponse.java b/src/main/java/de/extio/lmlib/client/oai/textcompletion/CompletionResponse.java index 66154ee..e659d3b 100755 --- a/src/main/java/de/extio/lmlib/client/oai/textcompletion/CompletionResponse.java +++ b/src/main/java/de/extio/lmlib/client/oai/textcompletion/CompletionResponse.java @@ -17,7 +17,9 @@ public final class CompletionResponse { private String content; private boolean isStoppedEos; - + + private Usage usage; + @JsonProperty("id") public String getId() { return this.id; @@ -71,7 +73,16 @@ public final class CompletionResponse { public void setStoppedEos(final boolean isStoppedEos) { this.isStoppedEos = isStoppedEos; } - + + @JsonProperty("usage") + public Usage getUsage() { + return usage; + } + + public void setUsage(final Usage usage) { + this.usage = usage; + } + @Override public String toString() { final StringBuilder builder = new StringBuilder(); @@ -87,6 +98,8 @@ public final class CompletionResponse { builder.append(this.content); builder.append(", isStoppedEos="); builder.append(this.isStoppedEos); + builder.append(", usage="); + builder.append(this.usage); builder.append("]"); return builder.toString(); } diff --git a/src/main/java/de/extio/lmlib/client/oai/textcompletion/TextCompletionClient.java b/src/main/java/de/extio/lmlib/client/oai/textcompletion/TextCompletionClient.java index 10a181c..e1a4234 100755 --- a/src/main/java/de/extio/lmlib/client/oai/textcompletion/TextCompletionClient.java +++ b/src/main/java/de/extio/lmlib/client/oai/textcompletion/TextCompletionClient.java @@ -162,6 +162,7 @@ public class TextCompletionClient implements Client, DisposableBean { final var response = webClient .method(HttpMethod.POST) .uri(uriBuilder -> uriBuilder.path("/v1/completions").build()) + .header("Content-Type", "application/json") .bodyValue(request) .retrieve() .bodyToMono(CompletionResponse.class) @@ -170,11 +171,7 @@ public class TextCompletionClient implements Client, DisposableBean { String content; CompletionFinishReason finishReason; - if (response.getContent() != null) { - content = promptStrategy.removeEOT(response.getContent()); - finishReason = response.isStoppedEos() ? CompletionFinishReason.DONE : CompletionFinishReason.TOKEN_LIMIT_REACHED; - } - else { + if (response.getChoices() != null && !response.getChoices().isEmpty()) { final var choice = response.getChoices().getFirst(); content = promptStrategy.removeEOT(choice.getText()); finishReason = switch (choice.getFinishReason()) { @@ -183,19 +180,35 @@ public class TextCompletionClient implements Client, DisposableBean { default -> CompletionFinishReason.DONE; }; } + else if (response.getContent() != null) { + content = promptStrategy.removeEOT(response.getContent()); + finishReason = response.isStoppedEos() ? CompletionFinishReason.DONE : CompletionFinishReason.TOKEN_LIMIT_REACHED; + } + else { + throw new IllegalStateException("No parseable response content from server"); + } + + int inTokens; + int outTokens; + if (response.getUsage() != null && response.getUsage().getPromptTokens() > 0) { + inTokens = response.getUsage().getPromptTokens(); + outTokens = response.getUsage().getCompletionTokens(); + } + else { + inTokens = this.tokenizer.count(prompt, modelProfile); + outTokens = this.tokenizer.count(content, modelProfile); + } - final var in = this.tokenizer.count(prompt, modelProfile); - final var out = this.tokenizer.count(content, modelProfile); final var statistics = new de.extio.lmlib.client.CompletionStatistics( 1, Duration.between(start, LocalDateTime.now()), - in, - out, - new BigDecimal(in).multiply(modelProfile.costInToken()).add(new BigDecimal(out).multiply(modelProfile.costOutToken())), + inTokens, + outTokens, + new BigDecimal(inTokens).multiply(modelProfile.costInToken()).add(new BigDecimal(outTokens).multiply(modelProfile.costOutToken())), false); LOGGER.debug("{}", statistics); if (this.collectStatistics) { - this.totalStatistics.add(statistics.duration(), in, out, statistics.cost()); + this.totalStatistics.add(statistics.duration(), inTokens, outTokens, statistics.cost()); } return new Completion(content, finishReason, statistics); diff --git a/src/main/java/de/extio/lmlib/client/oai/textcompletion/Usage.java b/src/main/java/de/extio/lmlib/client/oai/textcompletion/Usage.java new file mode 100644 index 0000000..f7f48df --- /dev/null +++ b/src/main/java/de/extio/lmlib/client/oai/textcompletion/Usage.java @@ -0,0 +1,45 @@ +package de.extio.lmlib.client.oai.textcompletion; + +import com.fasterxml.jackson.annotation.JsonProperty; + +public final class Usage { + + private int completionTokens; + + private int promptTokens; + + private int totalTokens; + + @JsonProperty("completion_tokens") + public int getCompletionTokens() { + return completionTokens; + } + + public void setCompletionTokens(final int completionTokens) { + this.completionTokens = completionTokens; + } + + @JsonProperty("prompt_tokens") + public int getPromptTokens() { + return promptTokens; + } + + public void setPromptTokens(final int promptTokens) { + this.promptTokens = promptTokens; + } + + @JsonProperty("total_tokens") + public int getTotalTokens() { + return totalTokens; + } + + public void setTotalTokens(final int totalTokens) { + this.totalTokens = totalTokens; + } + + @Override + public String toString() { + return "Usage [completionTokens=" + completionTokens + ", promptTokens=" + promptTokens + ", totalTokens=" + totalTokens + "]"; + } + +}