Skip to content

Commit e32efcf

Browse files
TheovanKraayTheovanKraay
committed
cosmos db entra id support and fixes
Signed-off-by: Theo van Kraay <[email protected]>
1 parent 5864255 commit e32efcf

File tree

8 files changed

+421
-37
lines changed

8 files changed

+421
-37
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

+25-6
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
/**
@@ -43,19 +45,37 @@
4345
* @author Soby Chacko
4446
* @since 1.0.0
4547
*/
48+
4649
@AutoConfiguration
4750
@ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
4851
@EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
4952
@ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
5053
matchIfMissing = true)
5154
public class CosmosDBVectorStoreAutoConfiguration {
5255

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

@@ -75,12 +95,11 @@ public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRe
7595
return CosmosDBVectorStore.builder(cosmosAsyncClient, embeddingModel)
7696
.databaseName(properties.getDatabaseName())
7797
.containerName(properties.getContainerName())
78-
.metadataFields(List.of(properties.getMetadataFields()))
98+
.metadataFields(properties.getMetadataFieldList())
7999
.vectorStoreThroughput(properties.getVectorStoreThroughput())
80100
.vectorDimensions(properties.getVectorDimensions())
81101
.partitionKeyPath(properties.getPartitionKeyPath())
82102
.build();
83-
84103
}
85104

86105
}

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

+19-1
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

+28-14
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);

‎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-vector-store</artifactId>

0 commit comments

Comments
 (0)