From fac38c5db86f48bb42156ad82b44241f088de195 Mon Sep 17 00:00:00 2001 From: Thomas Vitale Date: Fri, 12 Jul 2024 18:02:14 +0200 Subject: [PATCH] Update tests for EmbeddingResponseMetadata (#1039) Signed-off-by: Thomas Vitale --- .../ai/mistralai/MistralAiEmbeddingIT.java | 12 ++++++------ .../ai/openai/embedding/EmbeddingIT.java | 18 +++++++++--------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java index cd8197e1f..99d2aa1a6 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java @@ -40,9 +40,9 @@ class MistralAiEmbeddingIT { assertThat(embeddingResponse.getResults()).hasSize(1); assertThat(embeddingResponse.getResults().get(0)).isNotNull(); assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024); - assertThat(embeddingResponse.getMetadata()).containsEntry("model", "mistral-embed"); - assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 4); - assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 4); + assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("mistral-embed"); + assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(4); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(4); assertThat(mistralAiEmbeddingModel.dimensions()).isEqualTo(1024); } @@ -55,9 +55,9 @@ class MistralAiEmbeddingIT { assertThat(embeddingResponse.getResults()).hasSize(2); assertThat(embeddingResponse.getResults().get(0)).isNotNull(); assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1024); - assertThat(embeddingResponse.getMetadata()).containsEntry("model", "mistral-embed"); - assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 9); - assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 9); + assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("mistral-embed"); + assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(9); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(9); assertThat(mistralAiEmbeddingModel.dimensions()).isEqualTo(1024); } diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java index 326e9ebce..465fb5e70 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java @@ -42,9 +42,9 @@ class EmbeddingIT { assertThat(embeddingResponse.getResults()).hasSize(1); assertThat(embeddingResponse.getResults().get(0)).isNotNull(); assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1536); - assertThat(embeddingResponse.getMetadata()).containsEntry("model", "text-embedding-ada-002"); - assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2); - assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2); + assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("text-embedding-ada-002"); + assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(2); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(2); assertThat(embeddingModel.dimensions()).isEqualTo(1536); } @@ -57,9 +57,9 @@ class EmbeddingIT { assertThat(embeddingResponse.getResults()).hasSize(1); assertThat(embeddingResponse.getResults().get(0)).isNotNull(); assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(3072); - assertThat(embeddingResponse.getMetadata()).containsEntry("model", "text-embedding-3-large"); - assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2); - assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2); + assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("text-embedding-3-large"); + assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(2); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(2); // assertThat(embeddingModel.dimensions()).isEqualTo(3072); } @@ -73,9 +73,9 @@ class EmbeddingIT { assertThat(embeddingResponse.getResults().get(0)).isNotNull(); assertThat(embeddingResponse.getResults().get(0).getOutput()).hasSize(1536); - assertThat(embeddingResponse.getMetadata()).containsEntry("model", "text-embedding-3-small"); - assertThat(embeddingResponse.getMetadata()).containsEntry("total-tokens", 2); - assertThat(embeddingResponse.getMetadata()).containsEntry("prompt-tokens", 2); + assertThat(embeddingResponse.getMetadata().getModel()).isEqualTo("text-embedding-3-small"); + assertThat(embeddingResponse.getMetadata().getUsage().getTotalTokens()).isEqualTo(2); + assertThat(embeddingResponse.getMetadata().getUsage().getPromptTokens()).isEqualTo(2); // assertThat(embeddingModel.dimensions()).isEqualTo(3072); }