Fix NPE in OpenAiUsage

Fix gh-1152

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>

Author: ThomasVitale

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/metadata/OpenAiUsage.java

@@ -23,6 +23,7 @@
  * {@link Usage} implementation for {@literal OpenAI}.
  *
  * @author John Blum
+ * @author Thomas Vitale
  * @since 0.7.0
  * @see <a href=
  * "https://platform.openai.com/docs/api-reference/completions/object">Completion
@@ -47,17 +48,25 @@ protected OpenAiApi.Usage getUsage() {
 
 	@Override
 	public Long getPromptTokens() {
-		return getUsage().promptTokens().longValue();
+		Integer promptTokens = getUsage().promptTokens();
+		return promptTokens != null ? promptTokens.longValue() : 0;
 	}
 
 	@Override
 	public Long getGenerationTokens() {
-		return getUsage().completionTokens().longValue();
+		Integer generationTokens = getUsage().completionTokens();
+		return generationTokens != null ? generationTokens.longValue() : 0;
 	}
 
 	@Override
 	public Long getTotalTokens() {
-		return getUsage().totalTokens().longValue();
+		Integer totalTokens = getUsage().totalTokens();
+		if (totalTokens != null) {
+			return totalTokens.longValue();
+		}
+		else {
+			return getPromptTokens() + getGenerationTokens();
+		}
 	}
 
 	@Override

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/metadata/OpenAiUsageTests.java

@@ -0,0 +1,57 @@
+package org.springframework.ai.openai.metadata;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.openai.api.OpenAiApi;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests for {@link OpenAiUsage}.
+ *
+ * @author Thomas Vitale
+ */
+class OpenAiUsageTests {
+
+	@Test
+	void whenPromptTokensIsPresent() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getPromptTokens()).isEqualTo(200);
+	}
+
+	@Test
+	void whenPromptTokensIsNull() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, null, 100);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getPromptTokens()).isEqualTo(0);
+	}
+
+	@Test
+	void whenGenerationTokensIsPresent() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getGenerationTokens()).isEqualTo(100);
+	}
+
+	@Test
+	void whenGenerationTokensIsNull() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(null, 200, 200);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getGenerationTokens()).isEqualTo(0);
+	}
+
+	@Test
+	void whenTotalTokensIsPresent() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, 300);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getTotalTokens()).isEqualTo(300);
+	}
+
+	@Test
+	void whenTotalTokensIsNull() {
+		OpenAiApi.Usage openAiUsage = new OpenAiApi.Usage(100, 200, null);
+		OpenAiUsage usage = OpenAiUsage.from(openAiUsage);
+		assertThat(usage.getTotalTokens()).isEqualTo(300);
+	}
+
+}