|
39 | 39 | import org.springframework.ai.document.Document;
|
40 | 40 | import org.springframework.ai.embedding.EmbeddingModel;
|
41 | 41 | import org.springframework.ai.embedding.TokenCountBatchingStrategy;
|
42 |
| -import org.springframework.ai.mistralai.MistralAiEmbeddingModel; |
43 |
| -import org.springframework.ai.mistralai.api.MistralAiApi; |
44 | 42 | import org.springframework.ai.observation.conventions.SpringAiKind;
|
45 | 43 | import org.springframework.ai.observation.conventions.VectorStoreProvider;
|
| 44 | +import org.springframework.ai.openai.OpenAiEmbeddingModel; |
| 45 | +import org.springframework.ai.openai.api.OpenAiApi; |
46 | 46 | import org.springframework.ai.vectorstore.SearchRequest;
|
47 | 47 | import org.springframework.ai.vectorstore.VectorStore;
|
48 | 48 | import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
|
|
61 | 61 | * @author Thomas Vitale
|
62 | 62 | */
|
63 | 63 | @Testcontainers
|
64 |
| -@EnabledIfEnvironmentVariable(named = "MISTRAL_AI_API_KEY", matches = ".+") |
| 64 | +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") |
65 | 65 | public class QdrantVectorStoreObservationIT {
|
66 | 66 |
|
67 | 67 | private static final String COLLECTION_NAME = "test_collection";
|
68 | 68 |
|
69 |
| - private static final int EMBEDDING_DIMENSION = 1024; |
| 69 | + private static final int EMBEDDING_DIMENSION = 1536; |
70 | 70 |
|
71 | 71 | @Container
|
72 | 72 | static QdrantContainer qdrantContainer = new QdrantContainer(QdrantImage.DEFAULT_IMAGE);
|
@@ -126,7 +126,7 @@ void observationVectorStoreAddAndQueryOperations() {
|
126 | 126 | .hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(),
|
127 | 127 | SpringAiKind.VECTOR_STORE.value())
|
128 | 128 | .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString())
|
129 |
| - .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1024") |
| 129 | + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1536") |
130 | 130 | .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), COLLECTION_NAME)
|
131 | 131 | .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString())
|
132 | 132 | .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString())
|
@@ -159,7 +159,7 @@ void observationVectorStoreAddAndQueryOperations() {
|
159 | 159 |
|
160 | 160 | .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_QUERY_CONTENT.asString(),
|
161 | 161 | "What is Great Depression")
|
162 |
| - .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1024") |
| 162 | + .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_VECTOR_DIMENSION_COUNT.asString(), "1536") |
163 | 163 | .hasHighCardinalityKeyValue(HighCardinalityKeyNames.DB_COLLECTION_NAME.asString(), COLLECTION_NAME)
|
164 | 164 | .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_NAMESPACE.asString())
|
165 | 165 | .doesNotHaveHighCardinalityKeyValueWithKey(HighCardinalityKeyNames.DB_VECTOR_FIELD_NAME.asString())
|
@@ -206,7 +206,7 @@ public VectorStore qdrantVectorStore(EmbeddingModel embeddingModel, QdrantClient
|
206 | 206 |
|
207 | 207 | @Bean
|
208 | 208 | public EmbeddingModel embeddingModel() {
|
209 |
| - return new MistralAiEmbeddingModel(new MistralAiApi(System.getenv("MISTRAL_AI_API_KEY"))); |
| 209 | + return new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(System.getenv("OPENAI_API_KEY")).build()); |
210 | 210 | }
|
211 | 211 |
|
212 | 212 | }
|
|
0 commit comments