|
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