diff --git a/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponse.java b/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponse.java index 1f74f55..04e32e1 100644 --- a/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponse.java +++ b/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponse.java @@ -11,13 +11,16 @@ public class OpenAIChatCompletionResponse implements CompletionResponse { private final String id; private final List choices; + private final OpenAIChatCompletionResponseUsage usage; @JsonCreator(mode = JsonCreator.Mode.PROPERTIES) public OpenAIChatCompletionResponse( @JsonProperty("id") String id, - @JsonProperty("choices") List choices) { + @JsonProperty("choices") List choices, + @JsonProperty("usage") OpenAIChatCompletionResponseUsage usage) { this.id = id; this.choices = choices; + this.usage = usage; } public String getId() { @@ -27,4 +30,8 @@ public String getId() { public List getChoices() { return choices; } + + public OpenAIChatCompletionResponseUsage getUsage() { + return usage; + } } diff --git a/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponseUsage.java b/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponseUsage.java new file mode 100644 index 0000000..957fc17 --- /dev/null +++ b/src/main/java/ee/carlrobert/llm/client/openai/completion/response/OpenAIChatCompletionResponseUsage.java @@ -0,0 +1,35 @@ +package ee.carlrobert.llm.client.openai.completion.response; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; + +@JsonIgnoreProperties(ignoreUnknown = true) +public class OpenAIChatCompletionResponseUsage { + + private final int promptTokens; + private final int completionTokens; + private final int totalTokens; + + @JsonCreator(mode = JsonCreator.Mode.PROPERTIES) + public OpenAIChatCompletionResponseUsage( + @JsonProperty("prompt_tokens") int promptTokens, + @JsonProperty("completion_tokens") int completionTokens, + @JsonProperty("total_tokens") int totalTokens) { + this.promptTokens = promptTokens; + this.completionTokens = completionTokens; + this.totalTokens = totalTokens; + } + + public int getPromptTokens() { + return promptTokens; + } + + public int getCompletionTokens() { + return completionTokens; + } + + public int getTotalTokens() { + return totalTokens; + } +} \ No newline at end of file