Skip to content

Commit c979238

Browse files
tzolovtzolov
committed
Adding nav entry for Azure Cosmos DB documentation
- Minor code sytle fixes.
1 parent 7801119 commit c979238

File tree

10 files changed

+108
-77
lines changed

10 files changed

+108
-77
lines changed

‎spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
7676
* xref:api/vectordbs.adoc[]
7777
** xref:api/vectordbs/azure.adoc[]
78+
** xref:api/vectordbs/azure-cosmos-db.adoc[]
7879
** xref:api/vectordbs/apache-cassandra.adoc[]
7980
** xref:api/vectordbs/chroma.adoc[]
8081
** xref:api/vectordbs/elasticsearch.adoc[]

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,17 +101,19 @@ Add the following dependency to your Maven project:
101101

102102
The following configuration properties are available for the Cosmos DB vector store:
103103

104-
|===========================
104+
[stripes=even]
105+
|===
105106
| Property | Description
107+
106108
| spring.ai.vectorstore.cosmosdb.databaseName | The name of the Cosmos DB database to use.
107109
| spring.ai.vectorstore.cosmosdb.containerName | The name of the Cosmos DB container to use.
108110
| spring.ai.vectorstore.cosmosdb.partitionKeyPath | The path for the partition key.
109111
| spring.ai.vectorstore.cosmosdb.metadataFields | Comma-separated list of metadata fields.
110-
| spring.ai.vectorstore.cosmosdb.vectorStoreThoughput | The throughput for the vector store.
112+
| spring.ai.vectorstore.cosmosdb.vectorStoreThroughput | The throughput for the vector store.
111113
| spring.ai.vectorstore.cosmosdb.vectorDimensions | The number of dimensions for the vectors.
112114
| spring.ai.vectorstore.cosmosdb.endpoint | The endpoint for the Cosmos DB.
113115
| spring.ai.vectorstore.cosmosdb.key | The key for the Cosmos DB.
114-
|===========================
116+
|===
115117

116118

117119
== Complex Searches with Filters
@@ -205,7 +207,7 @@ public class DemoApplication implements CommandLineRunner {
205207
config.setDatabaseName("spring-ai-sample");
206208
config.setContainerName("container");
207209
config.setMetadataFields("country,city");
208-
config.setVectorStoreThoughput(400);
210+
config.setVectorStoreThroughput(400);
209211

210212
CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
211213
.endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfiguration.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRe
7373
config.setDatabaseName(properties.getDatabaseName());
7474
config.setContainerName(properties.getContainerName());
7575
config.setMetadataFields(properties.getMetadataFields());
76-
config.setVectorStoreThoughput(properties.getVectorStoreThoughput());
76+
config.setVectorStoreThroughput(properties.getVectorStoreThroughput());
7777
config.setVectorDimensions(properties.getVectorDimensions());
7878
return new CosmosDBVectorStore(observationRegistry, customObservationConvention.getIfAvailable(),
7979
cosmosAsyncClient, config, embeddingModel, batchingStrategy);

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreProperties.java

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
3535

3636
private String metadataFields;
3737

38-
private int vectorStoreThoughput = 400;
38+
private int vectorStoreThroughput = 400;
3939

4040
private long vectorDimensions = 1536;
4141

@@ -45,64 +45,64 @@ public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
4545

4646
private String key;
4747

48-
public int getVectorStoreThoughput() {
49-
return vectorStoreThoughput;
48+
public int getVectorStoreThroughput() {
49+
return this.vectorStoreThroughput;
5050
}
5151

52-
public void setVectorStoreThoughput(int vectorStoreThoughput) {
53-
this.vectorStoreThoughput = vectorStoreThoughput;
52+
public void setVectorStoreThroughput(int vectorStoreThroughput) {
53+
this.vectorStoreThroughput = vectorStoreThroughput;
5454
}
5555

5656
public String getMetadataFields() {
57-
return metadataFields;
57+
return this.metadataFields;
5858
}
5959

6060
public void setMetadataFields(String metadataFields) {
6161
this.metadataFields = metadataFields;
6262
}
6363

6464
public String getEndpoint() {
65-
return endpoint;
65+
return this.endpoint;
6666
}
6767

6868
public void setEndpoint(String endpoint) {
6969
this.endpoint = endpoint;
7070
}
7171

7272
public String getKey() {
73-
return key;
73+
return this.key;
7474
}
7575

7676
public void setKey(String key) {
7777
this.key = key;
7878
}
7979

8080
public String getDatabaseName() {
81-
return databaseName;
81+
return this.databaseName;
8282
}
8383

8484
public void setDatabaseName(String databaseName) {
8585
this.databaseName = databaseName;
8686
}
8787

8888
public String getContainerName() {
89-
return containerName;
89+
return this.containerName;
9090
}
9191

9292
public void setContainerName(String containerName) {
9393
this.containerName = containerName;
9494
}
9595

9696
public String getPartitionKeyPath() {
97-
return partitionKeyPath;
97+
return this.partitionKeyPath;
9898
}
9999

100100
public void setPartitionKeyPath(String partitionKeyPath) {
101101
this.partitionKeyPath = partitionKeyPath;
102102
}
103103

104104
public long getVectorDimensions() {
105-
return vectorDimensions;
105+
return this.vectorDimensions;
106106
}
107107

108108
public void setVectorDimensions(long vectorDimensions) {

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cosmosdb/CosmosDBVectorStoreAutoConfigurationIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public class CosmosDBVectorStoreAutoConfigurationIT {
5151
.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
5252
.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
5353
.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
54-
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThoughput=1000")
54+
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
5555
.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384")
5656
.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT"))
5757
.withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + System.getenv("AZURE_COSMOSDB_KEY"))

‎vector-stores/spring-ai-azure-cosmos-db-store/README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
[Azure Cosmos DB Vector Store Documentation]()
1+
[Azure Cosmos DB Vector Store Documentation](https://docs.spring.io/spring-ai/reference/api/vectordbs/azure-cosmos-db.html)

‎vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBFilterExpressionConverter.java

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,18 @@
1616

1717
package org.springframework.ai.vectorstore;
1818

19-
import org.springframework.ai.vectorstore.filter.Filter;
20-
import org.springframework.ai.vectorstore.filter.Filter.ExpressionType;
21-
import org.springframework.ai.vectorstore.filter.Filter.Key;
22-
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
19+
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
20+
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
2321

2422
import java.util.Collection;
2523
import java.util.Map;
2624
import java.util.Optional;
2725
import java.util.function.Function;
2826
import java.util.stream.Collectors;
2927

30-
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
31-
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
28+
import org.springframework.ai.vectorstore.filter.Filter;
29+
import org.springframework.ai.vectorstore.filter.Filter.Key;
30+
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
3231

3332
/**
3433
* Converts {@link org.springframework.ai.vectorstore.filter.Filter.Expression} into

‎vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/CosmosDBVectorStore.java

Lines changed: 67 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,15 @@
1616

1717
package org.springframework.ai.vectorstore;
1818

19-
import com.azure.cosmos.*;
20-
import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
21-
import com.azure.cosmos.models.*;
22-
import com.azure.cosmos.util.CosmosPagedFlux;
23-
import com.fasterxml.jackson.databind.JsonNode;
24-
import com.fasterxml.jackson.databind.ObjectMapper;
25-
import com.fasterxml.jackson.databind.node.ObjectNode;
26-
import io.micrometer.observation.ObservationRegistry;
19+
import java.util.ArrayList;
20+
import java.util.Collections;
21+
import java.util.HashMap;
22+
import java.util.List;
23+
import java.util.Optional;
24+
import java.util.stream.Collectors;
25+
import java.util.stream.IntStream;
26+
2727
import org.apache.commons.lang3.tuple.ImmutablePair;
28-
import org.apache.commons.lang3.tuple.Pair;
2928
import org.slf4j.Logger;
3029
import org.slf4j.LoggerFactory;
3130
import org.springframework.ai.document.Document;
@@ -38,10 +37,38 @@
3837
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
3938
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
4039
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
40+
41+
import com.azure.cosmos.CosmosAsyncClient;
42+
import com.azure.cosmos.CosmosAsyncContainer;
43+
import com.azure.cosmos.CosmosAsyncDatabase;
44+
import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
45+
import com.azure.cosmos.models.CosmosBulkOperations;
46+
import com.azure.cosmos.models.CosmosContainerProperties;
47+
import com.azure.cosmos.models.CosmosItemOperation;
48+
import com.azure.cosmos.models.CosmosQueryRequestOptions;
49+
import com.azure.cosmos.models.CosmosVectorDataType;
50+
import com.azure.cosmos.models.CosmosVectorDistanceFunction;
51+
import com.azure.cosmos.models.CosmosVectorEmbedding;
52+
import com.azure.cosmos.models.CosmosVectorEmbeddingPolicy;
53+
import com.azure.cosmos.models.CosmosVectorIndexSpec;
54+
import com.azure.cosmos.models.CosmosVectorIndexType;
55+
import com.azure.cosmos.models.ExcludedPath;
56+
import com.azure.cosmos.models.IncludedPath;
57+
import com.azure.cosmos.models.IndexingMode;
58+
import com.azure.cosmos.models.IndexingPolicy;
59+
import com.azure.cosmos.models.PartitionKey;
60+
import com.azure.cosmos.models.PartitionKeyDefinition;
61+
import com.azure.cosmos.models.PartitionKind;
62+
import com.azure.cosmos.models.SqlParameter;
63+
import com.azure.cosmos.models.SqlQuerySpec;
64+
import com.azure.cosmos.models.ThroughputProperties;
65+
import com.azure.cosmos.util.CosmosPagedFlux;
66+
import com.fasterxml.jackson.databind.JsonNode;
67+
import com.fasterxml.jackson.databind.ObjectMapper;
68+
import com.fasterxml.jackson.databind.node.ObjectNode;
69+
70+
import io.micrometer.observation.ObservationRegistry;
4171
import reactor.core.publisher.Flux;
42-
import java.util.*;
43-
import java.util.stream.Collectors;
44-
import java.util.stream.IntStream;
4572

4673
/**
4774
* @author Theo van Kraay
@@ -79,38 +106,38 @@ public CosmosDBVectorStore(ObservationRegistry observationRegistry,
79106
cosmosClient.createDatabaseIfNotExists(properties.getDatabaseName()).block();
80107

81108
initializeContainer(properties.getContainerName(), properties.getDatabaseName(),
82-
properties.getVectorStoreThoughput(), properties.getVectorDimensions(),
109+
properties.getVectorStoreThroughput(), properties.getVectorDimensions(),
83110
properties.getPartitionKeyPath());
84111

85112
this.embeddingModel = embeddingModel;
86113
}
87114

88-
private void initializeContainer(String containerName, String databaseName, int vectorStoreThoughput,
115+
private void initializeContainer(String containerName, String databaseName, int vectorStoreThroughput,
89116
long vectorDimensions, String partitionKeyPath) {
90117

91118
// Set defaults if not provided
92-
if (vectorStoreThoughput == 0) {
93-
vectorStoreThoughput = 400;
119+
if (vectorStoreThroughput == 0) {
120+
vectorStoreThroughput = 400;
94121
}
95122
if (partitionKeyPath == null) {
96123
partitionKeyPath = "/id";
97124
}
98125

99126
// handle hierarchical partition key
100-
PartitionKeyDefinition subpartitionKeyDefinition = new PartitionKeyDefinition();
101-
List<String> pathsfromCommaSeparatedList = new ArrayList<String>();
102-
String[] subpartitionKeyPaths = partitionKeyPath.split(",");
103-
Collections.addAll(pathsfromCommaSeparatedList, subpartitionKeyPaths);
104-
if (subpartitionKeyPaths.length > 1) {
105-
subpartitionKeyDefinition.setPaths(pathsfromCommaSeparatedList);
106-
subpartitionKeyDefinition.setKind(PartitionKind.MULTI_HASH);
127+
PartitionKeyDefinition subPartitionKeyDefinition = new PartitionKeyDefinition();
128+
List<String> pathsFromCommaSeparatedList = new ArrayList<String>();
129+
String[] subPartitionKeyPaths = partitionKeyPath.split(",");
130+
Collections.addAll(pathsFromCommaSeparatedList, subPartitionKeyPaths);
131+
if (subPartitionKeyPaths.length > 1) {
132+
subPartitionKeyDefinition.setPaths(pathsFromCommaSeparatedList);
133+
subPartitionKeyDefinition.setKind(PartitionKind.MULTI_HASH);
107134
}
108135
else {
109-
subpartitionKeyDefinition.setPaths(Collections.singletonList(partitionKeyPath));
110-
subpartitionKeyDefinition.setKind(PartitionKind.HASH);
136+
subPartitionKeyDefinition.setPaths(Collections.singletonList(partitionKeyPath));
137+
subPartitionKeyDefinition.setKind(PartitionKind.HASH);
111138
}
112139
CosmosContainerProperties collectionDefinition = new CosmosContainerProperties(containerName,
113-
subpartitionKeyDefinition);
140+
subPartitionKeyDefinition);
114141
// Set vector embedding policy
115142
CosmosVectorEmbeddingPolicy embeddingPolicy = new CosmosVectorEmbeddingPolicy();
116143
CosmosVectorEmbedding embedding = new CosmosVectorEmbedding();
@@ -135,16 +162,16 @@ private void initializeContainer(String containerName, String databaseName, int
135162
indexingPolicy.setVectorIndexes(List.of(cosmosVectorIndexSpec));
136163
collectionDefinition.setIndexingPolicy(indexingPolicy);
137164

138-
ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(vectorStoreThoughput);
139-
CosmosAsyncDatabase cosmosAsyncDatabase = cosmosClient.getDatabase(databaseName);
165+
ThroughputProperties throughputProperties = ThroughputProperties.createManualThroughput(vectorStoreThroughput);
166+
CosmosAsyncDatabase cosmosAsyncDatabase = this.cosmosClient.getDatabase(databaseName);
140167
cosmosAsyncDatabase.createContainerIfNotExists(collectionDefinition, throughputProperties).block();
141168
this.container = cosmosAsyncDatabase.getContainer(containerName);
142169
}
143170

144171
@Override
145172
public void close() {
146-
if (cosmosClient != null) {
147-
cosmosClient.close();
173+
if (this.cosmosClient != null) {
174+
this.cosmosClient.close();
148175
logger.info("Cosmos DB client closed successfully.");
149176
}
150177
}
@@ -192,7 +219,7 @@ public void doAdd(List<Document> documents) {
192219
.map(ImmutablePair::getValue)
193220
.collect(Collectors.toList());
194221

195-
container.executeBulkOperations(Flux.fromIterable(itemOperations)).doOnNext(response -> {
222+
this.container.executeBulkOperations(Flux.fromIterable(itemOperations)).doOnNext(response -> {
196223
if (response != null && response.getResponse() != null) {
197224
int statusCode = response.getResponse().getStatusCode();
198225
if (statusCode == 409) {
@@ -236,7 +263,7 @@ public Optional<Boolean> doDelete(List<String> idList) {
236263

237264
// Execute bulk delete operations synchronously by using blockLast() on the
238265
// Flux
239-
container.executeBulkOperations(Flux.fromIterable(itemOperations))
266+
this.container.executeBulkOperations(Flux.fromIterable(itemOperations))
240267
.doOnNext(response -> logger.info("Document deleted with status: {}",
241268
response.getResponse().getStatusCode()))
242269
.doOnError(error -> logger.error("Error deleting document: {}", error.getMessage()))
@@ -279,9 +306,11 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
279306
Filter.Expression filterExpression = request.getFilterExpression();
280307
if (filterExpression != null) {
281308
CosmosDBFilterExpressionConverter filterExpressionConverter = new CosmosDBFilterExpressionConverter(
282-
properties.getMetadataFieldsList()); // Use the expression directly as
283-
// it handles the "metadata"
284-
// fields internally
309+
this.properties.getMetadataFieldsList()); // Use the expression
310+
// directly as
311+
// it handles the
312+
// "metadata"
313+
// fields internally
285314
String filterQuery = filterExpressionConverter.convertExpression(filterExpression);
286315
queryBuilder.append(" AND ").append(filterQuery);
287316
}
@@ -297,7 +326,7 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
297326
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec(query, parameters);
298327
CosmosQueryRequestOptions options = new CosmosQueryRequestOptions();
299328

300-
CosmosPagedFlux<JsonNode> pagedFlux = container.queryItems(sqlQuerySpec, options, JsonNode.class);
329+
CosmosPagedFlux<JsonNode> pagedFlux = this.container.queryItems(sqlQuerySpec, options, JsonNode.class);
301330

302331
logger.info("Executing similarity search query: {}", query);
303332
try {
@@ -322,9 +351,9 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
322351
@Override
323352
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
324353
return VectorStoreObservationContext.builder(VectorStoreProvider.COSMOSDB.value(), operationName)
325-
.withCollectionName(container.getId())
354+
.withCollectionName(this.container.getId())
326355
.withDimensions(this.embeddingModel.dimensions())
327-
.withNamespace(container.getDatabase().getId())
356+
.withNamespace(this.container.getDatabase().getId())
328357
.withSimilarityMetric("cosine");
329358
}
330359

0 commit comments

Comments
 (0)