Skip to content

Commit 68b460a

Browse files
committed
Optimize MistralAiEmbeddingModel dimensions method
- Calculate and cache values for unknown models only if necessary - Make known embedding dimensions a mutable map attribute - Verify the cache mechanism with MistralAiEmbeddingModelTests Signed-off-by: Nicolas Krier <7557886+nicolaskrier@users.noreply.github.com>
1 parent 944029a commit 68b460a

File tree

1 file changed

+9
-2
lines changed

1 file changed

+9
-2
lines changed

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929

3030
import static org.assertj.core.api.Assertions.assertThat;
3131
import static org.mockito.ArgumentMatchers.any;
32+
import static org.mockito.Mockito.verify;
3233
import static org.mockito.Mockito.when;
3334

3435
/**
@@ -88,9 +89,15 @@ void testDimensionsFallbackForUnknownModel() {
8889
.retryTemplate(RetryUtils.DEFAULT_RETRY_TEMPLATE)
8990
.build();
9091

91-
// Should fall back to super.dimensions() which detects dimensions from the API
92-
// response
92+
// For the first call, it should fall back to super.dimensions() which detects
93+
// dimensions from the API response
9394
assertThat(model.dimensions()).isEqualTo(512);
95+
96+
// For the second call, it should use the cache mechanism.
97+
assertThat(model.dimensions()).isEqualTo(512);
98+
99+
// Verify that super.dimensions() has been called once.
100+
verify(mockApi).embeddings(any());
94101
}
95102

96103
@Test

0 commit comments

Comments
 (0)