Get token usage from OpenAI Api if available

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;

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

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

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 + "]";
+    }
+    
+}