Skip to content

Commit a9e1036

Browse files
tzolovtzolov
committed
clean code
1 parent 1000f46 commit a9e1036

19 files changed

+238
-364
lines changed

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vectorstore/weaviate/WeaviateVectorStoreAutoConfiguration.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package org.springframework.ai.autoconfigure.vectorstore.weaviate;
1717

18+
import io.micrometer.observation.ObservationRegistry;
1819
import io.weaviate.client.Config;
1920
import io.weaviate.client.WeaviateAuthClient;
2021
import io.weaviate.client.WeaviateClient;
@@ -23,6 +24,8 @@
2324
import org.springframework.ai.vectorstore.WeaviateVectorStore;
2425
import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig;
2526
import org.springframework.ai.vectorstore.WeaviateVectorStore.WeaviateVectorStoreConfig.MetadataField;
27+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
28+
import org.springframework.beans.factory.ObjectProvider;
2629
import org.springframework.boot.autoconfigure.AutoConfiguration;
2730
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2831
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -62,7 +65,8 @@ public WeaviateClient weaviateClient(WeaviateVectorStoreProperties properties,
6265
@Bean
6366
@ConditionalOnMissingBean
6467
public WeaviateVectorStore vectorStore(EmbeddingModel embeddingModel, WeaviateClient weaviateClient,
65-
WeaviateVectorStoreProperties properties) {
68+
WeaviateVectorStoreProperties properties, ObjectProvider<ObservationRegistry> observationRegistry,
69+
ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
6670

6771
WeaviateVectorStoreConfig.Builder configBuilder = WeaviateVectorStore.WeaviateVectorStoreConfig.builder()
6872
.withObjectClass(properties.getObjectClass())
@@ -73,7 +77,9 @@ public WeaviateVectorStore vectorStore(EmbeddingModel embeddingModel, WeaviateCl
7377
.toList())
7478
.withConsistencyLevel(properties.getConsistencyLevel());
7579

76-
return new WeaviateVectorStore(configBuilder.build(), embeddingModel, weaviateClient);
80+
return new WeaviateVectorStore(configBuilder.build(), embeddingModel, weaviateClient,
81+
observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
82+
customObservationConvention.getIfAvailable(() -> null));
7783
}
7884

7985
static class PropertiesWeaviateConnectionDetails implements WeaviateConnectionDetails {

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/azure/AzureVectorStoreAutoConfigurationIT.java‎

Lines changed: 13 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
*/
1616
package org.springframework.ai.autoconfigure.vectorstore.azure;
1717

18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.hamcrest.Matchers.hasSize;
20+
import static org.springframework.ai.autoconfigure.vectorstore.observation.ObservationTestUtil.assertObservationRegistry;
21+
1822
import java.io.IOException;
1923
import java.nio.charset.StandardCharsets;
2024
import java.time.Duration;
@@ -26,27 +30,21 @@
2630
import org.junit.jupiter.api.BeforeAll;
2731
import org.junit.jupiter.api.Test;
2832
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
29-
3033
import org.springframework.ai.document.Document;
3134
import org.springframework.ai.embedding.EmbeddingModel;
35+
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3236
import org.springframework.ai.transformers.TransformersEmbeddingModel;
33-
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
34-
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
35-
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames;
36-
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames;
3737
import org.springframework.ai.vectorstore.SearchRequest;
3838
import org.springframework.ai.vectorstore.VectorStore;
39+
import org.springframework.ai.vectorstore.azure.AzureVectorStore;
40+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
3941
import org.springframework.boot.autoconfigure.AutoConfigurations;
4042
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4143
import org.springframework.context.annotation.Bean;
4244
import org.springframework.context.annotation.Configuration;
4345
import org.springframework.core.io.DefaultResourceLoader;
4446

4547
import io.micrometer.observation.tck.TestObservationRegistry;
46-
import io.micrometer.observation.tck.TestObservationRegistryAssert;
47-
48-
import static org.assertj.core.api.Assertions.assertThat;
49-
import static org.hamcrest.Matchers.hasSize;
5048

5149
/**
5250
* @author Christian Tzolov
@@ -113,13 +111,8 @@ public void addAndSearchTest() {
113111
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
114112
}, hasSize(1));
115113

116-
TestObservationRegistryAssert.assertThat(observationRegistry)
117-
.doesNotHaveAnyRemainingCurrentObservation()
118-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
119-
.that()
120-
.hasContextualNameEqualTo("vector_store azure_vector_store add")
121-
.hasBeenStarted()
122-
.hasBeenStopped();
114+
assertObservationRegistry(observationRegistry, "vector_store", VectorStoreProvider.AZURE_VECTOR_STORE,
115+
VectorStoreObservationContext.Operation.ADD);
123116
observationRegistry.clear();
124117

125118
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
@@ -132,13 +125,8 @@ public void addAndSearchTest() {
132125
assertThat(resultDoc.getMetadata()).hasSize(2);
133126
assertThat(resultDoc.getMetadata()).containsKeys("spring", "distance");
134127

135-
TestObservationRegistryAssert.assertThat(observationRegistry)
136-
.doesNotHaveAnyRemainingCurrentObservation()
137-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
138-
.that()
139-
.hasContextualNameEqualTo("vector_store azure_vector_store query")
140-
.hasBeenStarted()
141-
.hasBeenStopped();
128+
assertObservationRegistry(observationRegistry, "vector_store", VectorStoreProvider.AZURE_VECTOR_STORE,
129+
VectorStoreObservationContext.Operation.QUERY);
142130
observationRegistry.clear();
143131

144132
// Remove all documents from the store
@@ -148,13 +136,8 @@ public void addAndSearchTest() {
148136
return vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
149137
}, hasSize(0));
150138

151-
TestObservationRegistryAssert.assertThat(observationRegistry)
152-
.doesNotHaveAnyRemainingCurrentObservation()
153-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
154-
.that()
155-
.hasContextualNameEqualTo("vector_store azure_vector_store delete")
156-
.hasBeenStarted()
157-
.hasBeenStopped();
139+
assertObservationRegistry(observationRegistry, "vector_store", VectorStoreProvider.AZURE_VECTOR_STORE,
140+
VectorStoreObservationContext.Operation.DELETE);
158141
observationRegistry.clear();
159142

160143
});

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/cassandra/CassandraVectorStoreAutoConfigurationIT.java‎

Lines changed: 16 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -15,32 +15,32 @@
1515
*/
1616
package org.springframework.ai.autoconfigure.vectorstore.cassandra;
1717

18+
import static org.assertj.core.api.Assertions.assertThat;
19+
import static org.springframework.ai.autoconfigure.vectorstore.observation.ObservationTestUtil.assertObservationRegistry;
20+
1821
import java.util.List;
1922
import java.util.Map;
2023

2124
import org.junit.jupiter.api.Test;
22-
import org.testcontainers.junit.jupiter.Container;
23-
import org.testcontainers.junit.jupiter.Testcontainers;
24-
import org.testcontainers.containers.CassandraContainer;
25-
import org.testcontainers.utility.DockerImageName;
26-
27-
import io.micrometer.observation.tck.TestObservationRegistry;
28-
import io.micrometer.observation.tck.TestObservationRegistryAssert;
29-
3025
import org.springframework.ai.ResourceUtils;
3126
import org.springframework.ai.document.Document;
3227
import org.springframework.ai.embedding.EmbeddingModel;
28+
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3329
import org.springframework.ai.transformers.TransformersEmbeddingModel;
3430
import org.springframework.ai.vectorstore.SearchRequest;
3531
import org.springframework.ai.vectorstore.VectorStore;
36-
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
32+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
3733
import org.springframework.boot.autoconfigure.AutoConfigurations;
3834
import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
3935
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4036
import org.springframework.context.annotation.Bean;
4137
import org.springframework.context.annotation.Configuration;
38+
import org.testcontainers.containers.CassandraContainer;
39+
import org.testcontainers.junit.jupiter.Container;
40+
import org.testcontainers.junit.jupiter.Testcontainers;
41+
import org.testcontainers.utility.DockerImageName;
4242

43-
import static org.assertj.core.api.Assertions.assertThat;
43+
import io.micrometer.observation.tck.TestObservationRegistry;
4444

4545
/**
4646
* @author Mick Semb Wever
@@ -80,13 +80,8 @@ void addAndSearch() {
8080
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
8181
vectorStore.add(documents);
8282

83-
TestObservationRegistryAssert.assertThat(observationRegistry)
84-
.doesNotHaveAnyRemainingCurrentObservation()
85-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
86-
.that()
87-
.hasContextualNameEqualTo("vector_store cassandra_vector_store add")
88-
.hasBeenStarted()
89-
.hasBeenStopped();
83+
assertObservationRegistry(observationRegistry, "vector_store",
84+
VectorStoreProvider.CASSANDRA_VECTOR_STORE, VectorStoreObservationContext.Operation.ADD);
9085
observationRegistry.clear();
9186

9287
List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
@@ -97,13 +92,8 @@ void addAndSearch() {
9792
assertThat(resultDoc.getContent()).contains(
9893
"Spring AI provides abstractions that serve as the foundation for developing AI applications.");
9994

100-
TestObservationRegistryAssert.assertThat(observationRegistry)
101-
.doesNotHaveAnyRemainingCurrentObservation()
102-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
103-
.that()
104-
.hasContextualNameEqualTo("vector_store cassandra_vector_store query")
105-
.hasBeenStarted()
106-
.hasBeenStopped();
95+
assertObservationRegistry(observationRegistry, "vector_store",
96+
VectorStoreProvider.CASSANDRA_VECTOR_STORE, VectorStoreObservationContext.Operation.QUERY);
10797
observationRegistry.clear();
10898

10999
// Remove all documents from the store
@@ -112,13 +102,8 @@ void addAndSearch() {
112102
results = vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(1));
113103
assertThat(results).isEmpty();
114104

115-
TestObservationRegistryAssert.assertThat(observationRegistry)
116-
.doesNotHaveAnyRemainingCurrentObservation()
117-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
118-
.that()
119-
.hasContextualNameEqualTo("vector_store cassandra_vector_store delete")
120-
.hasBeenStarted()
121-
.hasBeenStopped();
105+
assertObservationRegistry(observationRegistry, "vector_store",
106+
VectorStoreProvider.CASSANDRA_VECTOR_STORE, VectorStoreObservationContext.Operation.DELETE);
122107
observationRegistry.clear();
123108
});
124109
}

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/chroma/ChromaVectorStoreAutoConfigurationIT.java‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,20 @@
2828

2929
import org.springframework.ai.document.Document;
3030
import org.springframework.ai.embedding.EmbeddingModel;
31+
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3132
import org.springframework.ai.transformers.TransformersEmbeddingModel;
3233
import org.springframework.ai.vectorstore.SearchRequest;
3334
import org.springframework.ai.vectorstore.VectorStore;
3435
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
36+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
3537
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames;
3638
import org.springframework.boot.autoconfigure.AutoConfigurations;
3739
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
3840
import org.springframework.context.annotation.Bean;
3941
import org.springframework.context.annotation.Configuration;
4042

4143
import static org.assertj.core.api.Assertions.assertThat;
44+
import static org.springframework.ai.autoconfigure.vectorstore.observation.ObservationTestUtil.assertObservationRegistry;
4245

4346
/**
4447
* @author Christian Tzolov
@@ -74,13 +77,8 @@ public void addAndSearchWithFilters() {
7477

7578
vectorStore.add(List.of(bgDocument, nlDocument));
7679

77-
TestObservationRegistryAssert.assertThat(observationRegistry)
78-
.doesNotHaveAnyRemainingCurrentObservation()
79-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
80-
.that()
81-
.hasContextualNameEqualTo("vector_store chroma_vector_store add")
82-
.hasBeenStarted()
83-
.hasBeenStopped();
80+
assertObservationRegistry(observationRegistry, "vector_store", VectorStoreProvider.CHROMA_VECTOR_STORE,
81+
VectorStoreObservationContext.Operation.ADD);
8482
observationRegistry.clear();
8583

8684
var request = SearchRequest.query("The World").withTopK(5);

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/elasticsearch/ElasticsearchVectorStoreAutoConfigurationIT.java‎

Lines changed: 9 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import static org.assertj.core.api.Assertions.assertThat;
1919
import static org.hamcrest.Matchers.hasSize;
20+
import static org.springframework.ai.autoconfigure.vectorstore.observation.ObservationTestUtil.assertObservationRegistry;
2021

2122
import java.io.IOException;
2223
import java.nio.charset.StandardCharsets;
@@ -29,10 +30,11 @@
2930
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
3031
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
3132
import org.springframework.ai.document.Document;
33+
import org.springframework.ai.observation.conventions.VectorStoreProvider;
3234
import org.springframework.ai.vectorstore.ElasticsearchVectorStore;
3335
import org.springframework.ai.vectorstore.SearchRequest;
3436
import org.springframework.ai.vectorstore.SimilarityFunction;
35-
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
37+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
3638
import org.springframework.boot.autoconfigure.AutoConfigurations;
3739
import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
3840
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
@@ -45,7 +47,6 @@
4547
import org.testcontainers.junit.jupiter.Testcontainers;
4648

4749
import io.micrometer.observation.tck.TestObservationRegistry;
48-
import io.micrometer.observation.tck.TestObservationRegistryAssert;
4950

5051
@Testcontainers
5152
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -81,13 +82,8 @@ public void addAndSearchTest() {
8182

8283
vectorStore.add(documents);
8384

84-
TestObservationRegistryAssert.assertThat(observationRegistry)
85-
.doesNotHaveAnyRemainingCurrentObservation()
86-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
87-
.that()
88-
.hasContextualNameEqualTo("vector_store elasticsearch_vector_store add")
89-
.hasBeenStarted()
90-
.hasBeenStopped();
85+
assertObservationRegistry(observationRegistry, "vector_store",
86+
VectorStoreProvider.ELASTICSEARCH_VECTOR_STORE, VectorStoreObservationContext.Operation.ADD);
9187
observationRegistry.clear();
9288

9389
Awaitility.await()
@@ -108,25 +104,15 @@ public void addAndSearchTest() {
108104
assertThat(resultDoc.getMetadata()).containsKey("meta2");
109105
assertThat(resultDoc.getMetadata()).containsKey("distance");
110106

111-
TestObservationRegistryAssert.assertThat(observationRegistry)
112-
.doesNotHaveAnyRemainingCurrentObservation()
113-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
114-
.that()
115-
.hasContextualNameEqualTo("vector_store elasticsearch_vector_store query")
116-
.hasBeenStarted()
117-
.hasBeenStopped();
107+
assertObservationRegistry(observationRegistry, "vector_store",
108+
VectorStoreProvider.ELASTICSEARCH_VECTOR_STORE, VectorStoreObservationContext.Operation.QUERY);
118109
observationRegistry.clear();
119110

120111
// Remove all documents from the store
121112
vectorStore.delete(documents.stream().map(Document::getId).toList());
122113

123-
TestObservationRegistryAssert.assertThat(observationRegistry)
124-
.doesNotHaveAnyRemainingCurrentObservation()
125-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
126-
.that()
127-
.hasContextualNameEqualTo("vector_store elasticsearch_vector_store delete")
128-
.hasBeenStarted()
129-
.hasBeenStopped();
114+
assertObservationRegistry(observationRegistry, "vector_store",
115+
VectorStoreProvider.ELASTICSEARCH_VECTOR_STORE, VectorStoreObservationContext.Operation.DELETE);
130116
observationRegistry.clear();
131117

132118
Awaitility.await()

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/gemfire/GemFireVectorStoreAutoConfigurationIT.java‎

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,12 @@
1818

1919
import static org.assertj.core.api.Assertions.assertThat;
2020
import static org.hamcrest.Matchers.hasSize;
21+
import static org.springframework.ai.autoconfigure.vectorstore.observation.ObservationTestUtil.assertObservationRegistry;
2122

2223
import java.util.HashMap;
2324
import java.util.List;
2425
import java.util.Map;
2526

26-
import com.fasterxml.jackson.databind.JsonNode;
27-
import com.fasterxml.jackson.databind.ObjectMapper;
28-
import com.github.dockerjava.api.model.ExposedPort;
29-
import com.github.dockerjava.api.model.PortBinding;
30-
import com.github.dockerjava.api.model.Ports;
31-
import com.vmware.gemfire.testcontainers.GemFireCluster;
32-
33-
import io.micrometer.observation.tck.TestObservationRegistry;
34-
import io.micrometer.observation.tck.TestObservationRegistryAssert;
35-
3627
import org.awaitility.Awaitility;
3728
import org.junit.jupiter.api.AfterAll;
3829
import org.junit.jupiter.api.Assertions;
@@ -46,13 +37,21 @@
4637
import org.springframework.ai.vectorstore.GemFireVectorStore;
4738
import org.springframework.ai.vectorstore.SearchRequest;
4839
import org.springframework.ai.vectorstore.VectorStore;
49-
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
5040
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
5141
import org.springframework.boot.autoconfigure.AutoConfigurations;
5242
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
5343
import org.springframework.context.annotation.Bean;
5444
import org.springframework.context.annotation.Configuration;
5545

46+
import com.fasterxml.jackson.databind.JsonNode;
47+
import com.fasterxml.jackson.databind.ObjectMapper;
48+
import com.github.dockerjava.api.model.ExposedPort;
49+
import com.github.dockerjava.api.model.PortBinding;
50+
import com.github.dockerjava.api.model.Ports;
51+
import com.vmware.gemfire.testcontainers.GemFireCluster;
52+
53+
import io.micrometer.observation.tck.TestObservationRegistry;
54+
5655
/**
5756
* @author Geet Rawat
5857
* @author Christian Tzolov
@@ -225,15 +224,4 @@ public EmbeddingModel embeddingModel() {
225224

226225
}
227226

228-
private static void assertObservationRegistry(TestObservationRegistry observationRegistry, String kind,
229-
VectorStoreProvider vectorStoreProvider, VectorStoreObservationContext.Operation operation) {
230-
TestObservationRegistryAssert.assertThat(observationRegistry)
231-
.doesNotHaveAnyRemainingCurrentObservation()
232-
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
233-
.that()
234-
.hasContextualNameEqualTo(kind + " " + vectorStoreProvider.value() + " " + operation.value())
235-
.hasBeenStarted()
236-
.hasBeenStopped();
237-
}
238-
239227
}

0 commit comments

Comments
 (0)