Update tests for EmbeddingResponseMetadata (#1039)

Signed-off-by: Thomas Vitale <[email protected]>

Author: ThomasVitale

models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiEmbeddingIT.java

@@ -40,9 +40,9 @@ void defaultEmbedding() {
 		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 @@ void embeddingTest() {
 		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);
 	}
 

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/embedding/EmbeddingIT.java

@@ -42,9 +42,9 @@ void defaultEmbedding() {
 		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 @@ void embedding3Large() {
 		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 @@ void textEmbeddingAda002() {
 		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);
 	}