Skip to content

Commit e723371

Browse files
TheovanKraaysobychacko
TheovanKraay
authored and
sobychacko
committed
cosmos db entra id support and fixes
add doc changes Signed-off-by: Theo van Kraay <[email protected]>
1 parent c0bc623 commit e723371

File tree

9 files changed

+423
-42
lines changed

9 files changed

+423
-42
lines changed

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

Lines changed: 5 additions & 0 deletions
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‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,9 @@
1616

1717
package org.springframework.ai.vectorstore.cosmosdb.autoconfigure;
1818

19-
import java.util.List;
20-
2119
import com.azure.cosmos.CosmosAsyncClient;
2220
import com.azure.cosmos.CosmosClientBuilder;
21+
import com.azure.identity.DefaultAzureCredentialBuilder;
2322
import io.micrometer.observation.ObservationRegistry;
2423

2524
import org.springframework.ai.embedding.BatchingStrategy;
@@ -44,19 +43,37 @@
4443
* @author Soby Chacko
4544
* @since 1.0.0
4645
*/
46+
4747
@AutoConfiguration
4848
@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
4949
@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
5050
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
5151
matchIfMissing = true)
5252
public class CosmosDBVectorStoreAutoConfiguration {
5353

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

@@ -76,12 +93,11 @@ public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRe
7693
return CosmosDBVectorStore.builder(cosmosAsyncClient, embeddingModel)
7794
.databaseName(properties.getDatabaseName())
7895
.containerName(properties.getContainerName())
79-
.metadataFields(List.of(properties.getMetadataFields()))
96+
.metadataFields(properties.getMetadataFieldList())
8097
.vectorStoreThroughput(properties.getVectorStoreThroughput())
8198
.vectorDimensions(properties.getVectorDimensions())
8299
.partitionKeyPath(properties.getPartitionKeyPath())
83100
.build();
84-
85101
}
86102

87103
}

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

Lines changed: 19 additions & 1 deletion
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,12 @@ public void setMetadataFields(String metadataFields) {
6367
this.metadataFields = metadataFields;
6468
}
6569

70+
public List<String> getMetadataFieldList() {
71+
return this.metadataFields != null
72+
? Arrays.stream(this.metadataFields.split(",")).map(String::trim).filter(s -> !s.isEmpty()).toList()
73+
: List.of();
74+
}
75+
6676
public String getEndpoint() {
6777
return this.endpoint;
6878
}
@@ -79,6 +89,14 @@ public void setKey(String key) {
7989
this.key = key;
8090
}
8191

92+
public void setConnectionMode(String connectionMode) {
93+
this.connectionMode = connectionMode;
94+
}
95+
96+
public String getConnectionMode() {
97+
return this.connectionMode;
98+
}
99+
82100
public String getDatabaseName() {
83101
return this.databaseName;
84102
}

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

Lines changed: 28 additions & 14 deletions
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

@@ -124,14 +137,15 @@ 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()
136150
.query("The World")
137151
.topK(10)
@@ -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);

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ The following configuration properties are available for the Cosmos DB vector st
118118
| spring.ai.vectorstore.cosmosdb.vectorStoreThroughput | The throughput for the vector store.
119119
| spring.ai.vectorstore.cosmosdb.vectorDimensions | The number of dimensions for the vectors.
120120
| spring.ai.vectorstore.cosmosdb.endpoint | The endpoint for the Cosmos DB.
121-
| spring.ai.vectorstore.cosmosdb.key | The key for the Cosmos DB.
121+
| 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).
122122
|===
123123

124124

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

153153
== Setting up Azure Cosmos DB Vector Store without Auto Configuration
154154

155-
The following code demonstrates how to set up the `CosmosDBVectorStore` without relying on auto-configuration:
155+
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.
156156

157157
[source,java]
158158
----
@@ -161,7 +161,7 @@ public VectorStore vectorStore(ObservationRegistry observationRegistry) {
161161
// Create the Cosmos DB client
162162
CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
163163
.endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
164-
.key(System.getenv("COSMOSDB_AI_KEY"))
164+
.credential(new DefaultAzureCredentialBuilder().build())
165165
.userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
166166
.gatewayMode()
167167
.buildAsyncClient();

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

Lines changed: 5 additions & 0 deletions
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-vector-store</artifactId>

0 commit comments

Comments
 (0)