Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update tests for EmbeddingResponseMetadata (#1039)

Browse Source 
Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2024-07-12 18:02:14 +02:00
committed by GitHub
parent e601f14364
commit fac38c5db8
2 changed files with 15 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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