Skip to content

Commit 9e35834

Browse files
TheovanKraayTheovanKraay
committed
Signed-off-by: Theo van Kraay <[email protected]>
1 parent d813981 commit 9e35834

File tree

5 files changed

+122
-122
lines changed

5 files changed

+122
-122
lines changed

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/main/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfiguration.java

+20-20
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,8 @@
4848
@AutoConfiguration
4949
@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
5050
@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
51-
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB, matchIfMissing = true)
51+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
52+
matchIfMissing = true)
5253
public class CosmosDBVectorStoreAutoConfiguration {
5354

5455
private final String agentSuffix = "SpringAI-CDBNoSQL-VectorStore";
@@ -58,22 +59,23 @@ public CosmosAsyncClient cosmosClient(CosmosDBVectorStoreProperties properties)
5859
String mode = properties.getConnectionMode();
5960
if (mode == null) {
6061
properties.setConnectionMode("gateway");
61-
} else if (!mode.equals("direct") && !mode.equals("gateway")) {
62+
}
63+
else if (!mode.equals("direct") && !mode.equals("gateway")) {
6264
throw new IllegalArgumentException("Connection mode must be either 'direct' or 'gateway'");
6365
}
6466

65-
CosmosClientBuilder builder = new CosmosClientBuilder()
66-
.endpoint(properties.getEndpoint())
67-
.userAgentSuffix(agentSuffix);
67+
CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(properties.getEndpoint())
68+
.userAgentSuffix(agentSuffix);
6869

6970
if (properties.getKey() == null || properties.getKey().isEmpty()) {
7071
builder.credential(new DefaultAzureCredentialBuilder().build());
71-
} else {
72+
}
73+
else {
7274
builder.key(properties.getKey());
7375
}
7476

7577
return ("direct".equals(properties.getConnectionMode()) ? builder.directMode() : builder.gatewayMode())
76-
.buildAsyncClient();
78+
.buildAsyncClient();
7779
}
7880

7981
@Bean
@@ -84,21 +86,19 @@ BatchingStrategy batchingStrategy() {
8486

8587
@Bean
8688
@ConditionalOnMissingBean
87-
public CosmosDBVectorStore cosmosDBVectorStore(
88-
ObservationRegistry observationRegistry,
89+
public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRegistry,
8990
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
90-
CosmosDBVectorStoreProperties properties,
91-
CosmosAsyncClient cosmosAsyncClient,
92-
EmbeddingModel embeddingModel,
93-
BatchingStrategy batchingStrategy) {
91+
CosmosDBVectorStoreProperties properties, CosmosAsyncClient cosmosAsyncClient,
92+
EmbeddingModel embeddingModel, BatchingStrategy batchingStrategy) {
9493

9594
return CosmosDBVectorStore.builder(cosmosAsyncClient, embeddingModel)
96-
.databaseName(properties.getDatabaseName())
97-
.containerName(properties.getContainerName())
98-
.metadataFields(properties.getMetadataFieldList())
99-
.vectorStoreThroughput(properties.getVectorStoreThroughput())
100-
.vectorDimensions(properties.getVectorDimensions())
101-
.partitionKeyPath(properties.getPartitionKeyPath())
102-
.build();
95+
.databaseName(properties.getDatabaseName())
96+
.containerName(properties.getContainerName())
97+
.metadataFields(properties.getMetadataFieldList())
98+
.vectorStoreThroughput(properties.getVectorStoreThroughput())
99+
.vectorDimensions(properties.getVectorDimensions())
100+
.partitionKeyPath(properties.getPartitionKeyPath())
101+
.build();
103102
}
103+
104104
}

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/main/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreProperties.java

+3-4
Original file line numberDiff line numberDiff line change
@@ -68,10 +68,8 @@ public void setMetadataFields(String metadataFields) {
6868
}
6969

7070
public List<String> getMetadataFieldList() {
71-
return this.metadataFields != null ? Arrays.stream(this.metadataFields.split(","))
72-
.map(String::trim)
73-
.filter(s -> !s.isEmpty())
74-
.toList()
71+
return this.metadataFields != null
72+
? Arrays.stream(this.metadataFields.split(",")).map(String::trim).filter(s -> !s.isEmpty()).toList()
7573
: List.of();
7674
}
7775

@@ -130,4 +128,5 @@ public long getVectorDimensions() {
130128
public void setVectorDimensions(long vectorDimensions) {
131129
this.vectorDimensions = vectorDimensions;
132130
}
131+
133132
}

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/test/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfigurationIT.java

+25-23
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,13 @@ public CosmosDBVectorStoreAutoConfigurationIT() {
5555
String key = System.getenv("AZURE_COSMOSDB_KEY");
5656

5757
ApplicationContextRunner contextRunner = new ApplicationContextRunner()
58-
.withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class))
59-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database")
60-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
61-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
62-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
63-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
64-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384");
58+
.withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class))
59+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database")
60+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
61+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
62+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
63+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
64+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384");
6565

6666
if (endpoint != null && !"null".equalsIgnoreCase(endpoint)) {
6767
contextRunner = contextRunner.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + endpoint);
@@ -93,7 +93,7 @@ public void testAddSearchAndDeleteDocuments() {
9393

9494
// Perform a similarity search
9595
List<Document> results = this.vectorStore
96-
.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());
96+
.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());
9797

9898
// Verify the search results
9999
assertThat(results).isNotEmpty();
@@ -104,7 +104,7 @@ public void testAddSearchAndDeleteDocuments() {
104104

105105
// Perform a similarity search again
106106
List<Document> results2 = this.vectorStore
107-
.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());
107+
.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());
108108

109109
// Verify the search results
110110
assertThat(results2).isEmpty();
@@ -147,29 +147,29 @@ void testSimilaritySearchWithFilter() {
147147
FilterExpressionBuilder b = new FilterExpressionBuilder();
148148

149149
List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder()
150-
.query("The World")
151-
.topK(10)
152-
.filterExpression((b.in("country", "UK", "NL").build()))
153-
.build());
150+
.query("The World")
151+
.topK(10)
152+
.filterExpression((b.in("country", "UK", "NL").build()))
153+
.build());
154154

155155
assertThat(results).hasSize(2);
156156
assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
157157

158158
List<Document> results2 = this.vectorStore.similaritySearch(SearchRequest.builder()
159-
.query("The World")
160-
.topK(10)
161-
.filterExpression(
162-
b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build())
163-
.build());
159+
.query("The World")
160+
.topK(10)
161+
.filterExpression(
162+
b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build())
163+
.build());
164164

165165
assertThat(results2).hasSize(1);
166166
assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1");
167167

168168
List<Document> results3 = this.vectorStore.similaritySearch(SearchRequest.builder()
169-
.query("The World")
170-
.topK(10)
171-
.filterExpression(b.and(b.eq("country", "US"), b.eq("year", 2020)).build())
172-
.build());
169+
.query("The World")
170+
.topK(10)
171+
.filterExpression(b.and(b.eq("country", "US"), b.eq("year", 2020)).build())
172+
.build());
173173

174174
assertThat(results3).hasSize(1);
175175
assertThat(results3).extracting(Document::getId).containsExactlyInAnyOrder("4");
@@ -178,7 +178,7 @@ void testSimilaritySearchWithFilter() {
178178

179179
// Perform a similarity search again
180180
List<Document> results4 = this.vectorStore
181-
.similaritySearch(SearchRequest.builder().query("The World").topK(1).build());
181+
.similaritySearch(SearchRequest.builder().query("The World").topK(1).build());
182182

183183
// Verify the search results
184184
assertThat(results4).isEmpty();
@@ -223,5 +223,7 @@ public EmbeddingModel embeddingModel() {
223223
public TestObservationRegistry observationRegistry() {
224224
return TestObservationRegistry.create();
225225
}
226+
226227
}
228+
227229
}

0 commit comments

Comments
 (0)