Skip to content

Commit 079e512

Browse files
TheovanKraayTheovanKraay
committed
Entra ID support for auto config
Signed-off-by: Theo van Kraay <[email protected]>
1 parent 3ca8d70 commit 079e512

File tree

8 files changed

+199
-131
lines changed

8 files changed

+199
-131
lines changed

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/pom.xml

+5
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,11 @@
4545
<version>${project.parent.version}</version>
4646
<optional>true</optional>
4747
</dependency>
48+
<dependency>
49+
<groupId>com.azure</groupId>
50+
<artifactId>azure-identity</artifactId>
51+
<version>1.15.4</version> <!-- or the latest version -->
52+
</dependency>
4853
<dependency>
4954
<groupId>org.springframework.boot</groupId>
5055
<artifactId>spring-boot-starter</artifactId>

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

+37-22
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import com.azure.cosmos.CosmosAsyncClient;
2020
import com.azure.cosmos.CosmosClientBuilder;
21+
import com.azure.identity.DefaultAzureCredentialBuilder;
2122
import io.micrometer.observation.ObservationRegistry;
2223

2324
import org.springframework.ai.embedding.BatchingStrategy;
@@ -33,6 +34,7 @@
3334
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3435
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3536
import org.springframework.context.annotation.Bean;
37+
3638
import java.util.List;
3739

3840
/**
@@ -42,24 +44,36 @@
4244
* @author Soby Chacko
4345
* @since 1.0.0
4446
*/
47+
4548
@AutoConfiguration
4649
@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
4750
@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
48-
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
49-
matchIfMissing = true)
51+
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB, matchIfMissing = true)
5052
public class CosmosDBVectorStoreAutoConfiguration {
5153

52-
String endpoint;
53-
54-
String key;
54+
private final String agentSuffix = "SpringAI-CDBNoSQL-VectorStore";
5555

5656
@Bean
5757
public CosmosAsyncClient cosmosClient(CosmosDBVectorStoreProperties properties) {
58-
return new CosmosClientBuilder().endpoint(properties.getEndpoint())
59-
.userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
60-
.key(properties.getKey())
61-
.gatewayMode()
62-
.buildAsyncClient();
58+
String mode = properties.getConnectionMode();
59+
if (mode == null) {
60+
properties.setConnectionMode("gateway");
61+
} else if (!mode.equals("direct") && !mode.equals("gateway")) {
62+
throw new IllegalArgumentException("Connection mode must be either 'direct' or 'gateway'");
63+
}
64+
65+
CosmosClientBuilder builder = new CosmosClientBuilder()
66+
.endpoint(properties.getEndpoint())
67+
.userAgentSuffix(agentSuffix);
68+
69+
if (properties.getKey() == null || properties.getKey().isEmpty()) {
70+
builder.credential(new DefaultAzureCredentialBuilder().build());
71+
} else {
72+
builder.key(properties.getKey());
73+
}
74+
75+
return ("direct".equals(properties.getConnectionMode()) ? builder.directMode() : builder.gatewayMode())
76+
.buildAsyncClient();
6377
}
6478

6579
@Bean
@@ -70,20 +84,21 @@ BatchingStrategy batchingStrategy() {
7084

7185
@Bean
7286
@ConditionalOnMissingBean
73-
public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRegistry,
87+
public CosmosDBVectorStore cosmosDBVectorStore(
88+
ObservationRegistry observationRegistry,
7489
ObjectProvider<VectorStoreObservationConvention> customObservationConvention,
75-
CosmosDBVectorStoreProperties properties, CosmosAsyncClient cosmosAsyncClient,
76-
EmbeddingModel embeddingModel, BatchingStrategy batchingStrategy) {
90+
CosmosDBVectorStoreProperties properties,
91+
CosmosAsyncClient cosmosAsyncClient,
92+
EmbeddingModel embeddingModel,
93+
BatchingStrategy batchingStrategy) {
7794

7895
return CosmosDBVectorStore.builder(cosmosAsyncClient, embeddingModel)
79-
.databaseName(properties.getDatabaseName())
80-
.containerName(properties.getContainerName())
81-
.metadataFields(List.of(properties.getMetadataFields()))
82-
.vectorStoreThroughput(properties.getVectorStoreThroughput())
83-
.vectorDimensions(properties.getVectorDimensions())
84-
.partitionKeyPath(properties.getPartitionKeyPath())
85-
.build();
86-
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();
87103
}
88-
89104
}

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

+21-2
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,15 @@
1919
import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties;
2020
import org.springframework.boot.context.properties.ConfigurationProperties;
2121

22+
import java.util.Arrays;
23+
import java.util.List;
24+
2225
/**
2326
* Configuration properties for CosmosDB Vector Store.
2427
*
2528
* @author Theo van Kraay
2629
* @since 1.0.0
2730
*/
28-
2931
@ConfigurationProperties(CosmosDBVectorStoreProperties.CONFIG_PREFIX)
3032
public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
3133

@@ -47,6 +49,8 @@ public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
4749

4850
private String key;
4951

52+
private String connectionMode;
53+
5054
public int getVectorStoreThroughput() {
5155
return this.vectorStoreThroughput;
5256
}
@@ -63,6 +67,14 @@ public void setMetadataFields(String metadataFields) {
6367
this.metadataFields = metadataFields;
6468
}
6569

70+
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()
75+
: List.of();
76+
}
77+
6678
public String getEndpoint() {
6779
return this.endpoint;
6880
}
@@ -79,6 +91,14 @@ public void setKey(String key) {
7991
this.key = key;
8092
}
8193

94+
public void setConnectionMode(String connectionMode) {
95+
this.connectionMode = connectionMode;
96+
}
97+
98+
public String getConnectionMode() {
99+
return this.connectionMode;
100+
}
101+
82102
public String getDatabaseName() {
83103
return this.databaseName;
84104
}
@@ -110,5 +130,4 @@ public long getVectorDimensions() {
110130
public void setVectorDimensions(long vectorDimensions) {
111131
this.vectorDimensions = vectorDimensions;
112132
}
113-
114133
}

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

+44-32
Original file line numberDiff line numberDiff line change
@@ -44,22 +44,35 @@
4444
* @author Theo van Kraay
4545
* @since 1.0.0
4646
*/
47-
4847
@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+")
4948
@EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_KEY", matches = ".+")
5049
public class CosmosDBVectorStoreAutoConfigurationIT {
5150

52-
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
53-
.withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class))
54-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database")
55-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
56-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
57-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
58-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
59-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384")
60-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT"))
61-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + System.getenv("AZURE_COSMOSDB_KEY"))
62-
.withUserConfiguration(Config.class);
51+
private final ApplicationContextRunner contextRunner;
52+
53+
public CosmosDBVectorStoreAutoConfigurationIT() {
54+
String endpoint = System.getenv("AZURE_COSMOSDB_ENDPOINT");
55+
String key = System.getenv("AZURE_COSMOSDB_KEY");
56+
57+
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");
65+
66+
if (endpoint != null && !"null".equalsIgnoreCase(endpoint)) {
67+
contextRunner = contextRunner.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + endpoint);
68+
}
69+
70+
if (key != null && !"null".equalsIgnoreCase(key)) {
71+
contextRunner = contextRunner.withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + key);
72+
}
73+
74+
this.contextRunner = contextRunner.withUserConfiguration(Config.class);
75+
}
6376

6477
private VectorStore vectorStore;
6578

@@ -80,7 +93,7 @@ public void testAddSearchAndDeleteDocuments() {
8093

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

8598
// Verify the search results
8699
assertThat(results).isNotEmpty();
@@ -91,7 +104,7 @@ public void testAddSearchAndDeleteDocuments() {
91104

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

96109
// Verify the search results
97110
assertThat(results2).isEmpty();
@@ -124,38 +137,39 @@ void testSimilaritySearchWithFilter() {
124137
metadata4.put("country", "US");
125138
metadata4.put("year", 2020);
126139
metadata4.put("city", "Sofia");
127-
128140
Document document1 = new Document("1", "A document about the UK", metadata1);
129141
Document document2 = new Document("2", "A document about the Netherlands", metadata2);
130142
Document document3 = new Document("3", "A document about the US", metadata3);
131143
Document document4 = new Document("4", "A document about the US", metadata4);
132144

133145
this.vectorStore.add(List.of(document1, document2, document3, document4));
146+
134147
FilterExpressionBuilder b = new FilterExpressionBuilder();
148+
135149
List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder()
136-
.query("The World")
137-
.topK(10)
138-
.filterExpression((b.in("country", "UK", "NL").build()))
139-
.build());
150+
.query("The World")
151+
.topK(10)
152+
.filterExpression((b.in("country", "UK", "NL").build()))
153+
.build());
140154

141155
assertThat(results).hasSize(2);
142156
assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
143157

144158
List<Document> results2 = this.vectorStore.similaritySearch(SearchRequest.builder()
145-
.query("The World")
146-
.topK(10)
147-
.filterExpression(
148-
b.and(b.or(b.gte("year", 2021), b.eq("country", "NL")), b.ne("city", "Amsterdam")).build())
149-
.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());
150164

151165
assertThat(results2).hasSize(1);
152166
assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1");
153167

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

160174
assertThat(results3).hasSize(1);
161175
assertThat(results3).extracting(Document::getId).containsExactlyInAnyOrder("4");
@@ -164,7 +178,7 @@ void testSimilaritySearchWithFilter() {
164178

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

169183
// Verify the search results
170184
assertThat(results4).isEmpty();
@@ -190,7 +204,7 @@ public void autoConfigurationEnabledByDefault() {
190204

191205
@Test
192206
public void autoConfigurationEnabledWhenTypeIsAzureCosmosDB() {
193-
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure-cosmmos-db").run(context -> {
207+
this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure-cosmos-db").run(context -> {
194208
assertThat(context.getBeansOfType(CosmosDBVectorStoreProperties.class)).isNotEmpty();
195209
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
196210
assertThat(context.getBean(VectorStore.class)).isInstanceOf(CosmosDBVectorStore.class);
@@ -209,7 +223,5 @@ public EmbeddingModel embeddingModel() {
209223
public TestObservationRegistry observationRegistry() {
210224
return TestObservationRegistry.create();
211225
}
212-
213226
}
214-
215227
}

‎spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc

+3-3
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ The following configuration properties are available for the Cosmos DB vector st
112112
| spring.ai.vectorstore.cosmosdb.vectorStoreThroughput | The throughput for the vector store.
113113
| spring.ai.vectorstore.cosmosdb.vectorDimensions | The number of dimensions for the vectors.
114114
| spring.ai.vectorstore.cosmosdb.endpoint | The endpoint for the Cosmos DB.
115-
| spring.ai.vectorstore.cosmosdb.key | The key for the Cosmos DB.
115+
| spring.ai.vectorstore.cosmosdb.key | The key for the Cosmos DB (if key is not present, [DefaultAzureCredential](https://learn.microsoft.com/azure/developer/java/sdk/authentication/credential-chains#defaultazurecredential-overview) will be used).
116116
|===
117117

118118

@@ -146,7 +146,7 @@ List<Document> results = vectorStore.similaritySearch(SearchRequest.builder().qu
146146

147147
== Setting up Azure Cosmos DB Vector Store without Auto Configuration
148148

149-
The following code demonstrates how to set up the `CosmosDBVectorStore` without relying on auto-configuration:
149+
The following code demonstrates how to set up the `CosmosDBVectorStore` without relying on auto-configuration. [DefaultAzureCredential](https://learn.microsoft.com/azure/developer/java/sdk/authentication/credential-chains#defaultazurecredential-overview) is recommended for authentication to Azure Cosmos DB.
150150

151151
[source,java]
152152
----
@@ -155,7 +155,7 @@ public VectorStore vectorStore(ObservationRegistry observationRegistry) {
155155
// Create the Cosmos DB client
156156
CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
157157
.endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
158-
.key(System.getenv("COSMOSDB_AI_KEY"))
158+
.credential(new DefaultAzureCredentialBuilder().build())
159159
.userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
160160
.gatewayMode()
161161
.buildAsyncClient();

‎vector-stores/spring-ai-azure-cosmos-db-store/pom.xml

+5
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@
4747
<artifactId>azure-spring-data-cosmos</artifactId>
4848
<version>${azure-cosmos.version}</version>
4949
</dependency>
50+
<dependency>
51+
<groupId>com.azure</groupId>
52+
<artifactId>azure-identity</artifactId>
53+
<version>1.15.4</version> <!-- or the latest version -->
54+
</dependency>
5055
<dependency>
5156
<groupId>org.springframework.ai</groupId>
5257
<artifactId>spring-ai-core</artifactId>

0 commit comments

Comments
 (0)