Skip to content

Commit 7a8cdab

Browse files
sobychackosobychacko
committed
Introducing a Base Builder hierarchy for VectorStore
 - Trying the builder hierarchy on Chroma vector store
1 parent f54b3d5 commit 7a8cdab

File tree

11 files changed

+326
-105
lines changed

11 files changed

+326
-105
lines changed

‎spring-ai-core/src/main/java/org/springframework/ai/vectorstore/AbstractVectorStoreBuilder.java‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.vectorstore;
18+
19+
import io.micrometer.observation.ObservationRegistry;
20+
21+
import org.springframework.ai.embedding.EmbeddingModel;
22+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
23+
import org.springframework.lang.Nullable;
24+
import org.springframework.util.Assert;
25+
26+
/**
27+
* Abstract base builder implementing common builder functionality for
28+
* {@link VectorStore}. Provides default implementations for observation-related settings.
29+
*
30+
* @param <T> the concrete builder type, enabling method chaining with the correct return
31+
* type
32+
*/
33+
public abstract class AbstractVectorStoreBuilder<T extends AbstractVectorStoreBuilder<T>>
34+
implements VectorStore.Builder<T> {
35+
36+
private final EmbeddingModel embeddingModel;
37+
38+
protected ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
39+
40+
@Nullable
41+
protected VectorStoreObservationConvention customObservationConvention;
42+
43+
/**
44+
* Creates a new builder instance with the required embedding model.
45+
* @param embeddingModel the embedding model to use for vector operations
46+
* @throws IllegalArgumentException if embeddingModel is null
47+
*/
48+
protected AbstractVectorStoreBuilder(EmbeddingModel embeddingModel) {
49+
Assert.notNull(embeddingModel, "EmbeddingModel must not be null");
50+
this.embeddingModel = embeddingModel;
51+
}
52+
53+
/**
54+
* Returns this builder cast to the concrete builder type. Used internally to enable
55+
* proper method chaining in subclasses.
56+
* @return this builder cast to the concrete type
57+
*/
58+
@SuppressWarnings("unchecked")
59+
protected T self() {
60+
return (T) this;
61+
}
62+
63+
@Override
64+
public T observationRegistry(ObservationRegistry observationRegistry) {
65+
this.observationRegistry = observationRegistry;
66+
return self();
67+
}
68+
69+
@Override
70+
public T customObservationConvention(VectorStoreObservationConvention convention) {
71+
this.customObservationConvention = convention;
72+
return self();
73+
}
74+
75+
/**
76+
* Returns the configured embedding model.
77+
* @return the embedding model
78+
*/
79+
@Override
80+
public EmbeddingModel embeddingModel() {
81+
return this.embeddingModel;
82+
}
83+
84+
/**
85+
* Returns the configured observation registry.
86+
* @return the observation registry, never null
87+
*/
88+
@Override
89+
public ObservationRegistry observationRegistry() {
90+
return this.observationRegistry;
91+
}
92+
93+
/**
94+
* Returns the configured custom observation convention.
95+
* @return the custom observation convention, may be null
96+
*/
97+
@Override
98+
@Nullable
99+
public VectorStoreObservationConvention customObservationConvention() {
100+
return this.customObservationConvention;
101+
}
102+
103+
}

‎spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,14 @@
1919
import java.util.List;
2020
import java.util.Optional;
2121

22+
import io.micrometer.observation.ObservationRegistry;
23+
2224
import org.springframework.ai.document.Document;
2325
import org.springframework.ai.document.DocumentWriter;
26+
import org.springframework.ai.embedding.EmbeddingModel;
27+
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
28+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
29+
import org.springframework.lang.Nullable;
2430

2531
/**
2632
* The {@code VectorStore} interface defines the operations for managing and querying
@@ -75,4 +81,56 @@ default List<Document> similaritySearch(String query) {
7581
return this.similaritySearch(SearchRequest.query(query));
7682
}
7783

84+
/**
85+
* Builder interface for creating VectorStore instances. Implements a fluent builder
86+
* pattern for configuring observation-related settings.
87+
*
88+
* @param <T> the concrete builder type, enabling method chaining with the correct
89+
* return type
90+
*/
91+
interface Builder<T extends Builder<T>> {
92+
93+
/**
94+
* Sets the registry for collecting observations and metrics. Defaults to
95+
* {@link ObservationRegistry#NOOP} if not specified.
96+
* @param observationRegistry the registry to use for observations
97+
* @return the builder instance for method chaining
98+
*/
99+
T observationRegistry(ObservationRegistry observationRegistry);
100+
101+
/**
102+
* Sets a custom convention for creating observations. If not specified,
103+
* {@link DefaultVectorStoreObservationConvention} will be used.
104+
* @param convention the custom observation convention to use
105+
* @return the builder instance for method chaining
106+
*/
107+
T customObservationConvention(VectorStoreObservationConvention convention);
108+
109+
/**
110+
* Returns the configured embedding model.
111+
* @return the embedding model, never null
112+
*/
113+
EmbeddingModel embeddingModel();
114+
115+
/**
116+
* Returns the configured observation registry.
117+
* @return the observation registry, never null
118+
*/
119+
ObservationRegistry observationRegistry();
120+
121+
/**
122+
* Returns the configured custom observation convention.
123+
* @return the custom observation convention, may be null
124+
*/
125+
@Nullable
126+
VectorStoreObservationConvention customObservationConvention();
127+
128+
/**
129+
* Builds and returns a new VectorStore instance with the configured settings.
130+
* @return a new VectorStore instance
131+
*/
132+
VectorStore build();
133+
134+
}
135+
78136
}

‎spring-ai-core/src/main/java/org/springframework/ai/vectorstore/observation/NewAbstractObservationVectorStore.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
/*
2+
* Copyright 2023-2024 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.vectorstore.observation;
18+
19+
import java.util.List;
20+
import java.util.Optional;
21+
22+
import io.micrometer.observation.ObservationRegistry;
23+
24+
import org.springframework.ai.document.Document;
25+
import org.springframework.ai.embedding.EmbeddingModel;
26+
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
27+
import org.springframework.ai.vectorstore.SearchRequest;
28+
import org.springframework.ai.vectorstore.VectorStore;
29+
import org.springframework.lang.Nullable;
30+
31+
/**
32+
* NOTE: This class is a temporary class. Once the review is done and if we decide to go
33+
* ahead, we will move this class to AbstractObservationVectorStore.
34+
*
35+
* Abstract base implementation of VectorStore that provides observation capabilities.
36+
* Handles common observation-related functionality for concrete vector store
37+
* implementations.
38+
*/
39+
public abstract class NewAbstractObservationVectorStore implements VectorStore {
40+
41+
private static final VectorStoreObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultVectorStoreObservationConvention();
42+
43+
protected final EmbeddingModel embeddingModel;
44+
45+
private final ObservationRegistry observationRegistry;
46+
47+
@Nullable
48+
private final VectorStoreObservationConvention customObservationConvention;
49+
50+
/**
51+
* Creates a new AbstractObservationVectorStore instance with the specified builder
52+
* settings. Initializes observation-related components and the embedding model.
53+
* @param builder the builder containing configuration settings
54+
*/
55+
public NewAbstractObservationVectorStore(AbstractVectorStoreBuilder<?> builder) {
56+
this.embeddingModel = builder.embeddingModel();
57+
this.observationRegistry = builder.observationRegistry();
58+
this.customObservationConvention = builder.customObservationConvention();
59+
}
60+
61+
/**
62+
* Create a new {@link AbstractObservationVectorStore} instance.
63+
* @param documents the documents to add
64+
*/
65+
@Override
66+
public void add(List<Document> documents) {
67+
68+
VectorStoreObservationContext observationContext = this
69+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.ADD.value())
70+
.build();
71+
72+
VectorStoreObservationDocumentation.AI_VECTOR_STORE
73+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
74+
this.observationRegistry)
75+
.observe(() -> this.doAdd(documents));
76+
}
77+
78+
@Override
79+
public Optional<Boolean> delete(List<String> deleteDocIds) {
80+
81+
VectorStoreObservationContext observationContext = this
82+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value())
83+
.build();
84+
85+
return VectorStoreObservationDocumentation.AI_VECTOR_STORE
86+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
87+
this.observationRegistry)
88+
.observe(() -> this.doDelete(deleteDocIds));
89+
}
90+
91+
@Override
92+
public List<Document> similaritySearch(SearchRequest request) {
93+
94+
VectorStoreObservationContext searchObservationContext = this
95+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
96+
.withQueryRequest(request)
97+
.build();
98+
99+
return VectorStoreObservationDocumentation.AI_VECTOR_STORE
100+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION,
101+
() -> searchObservationContext, this.observationRegistry)
102+
.observe(() -> {
103+
var documents = this.doSimilaritySearch(request);
104+
searchObservationContext.setQueryResponse(documents);
105+
return documents;
106+
});
107+
}
108+
109+
/**
110+
* Perform the actual add operation.
111+
* @param documents the documents to add
112+
*/
113+
public abstract void doAdd(List<Document> documents);
114+
115+
/**
116+
* Perform the actual delete operation.
117+
* @param idList the list of document IDs to delete
118+
* @return true if the documents were successfully deleted
119+
*/
120+
public abstract Optional<Boolean> doDelete(List<String> idList);
121+
122+
/**
123+
* Perform the actual similarity search operation.
124+
* @param request the search request
125+
* @return the list of documents that match the query request conditions
126+
*/
127+
public abstract List<Document> doSimilaritySearch(SearchRequest request);
128+
129+
/**
130+
* Create a new {@link VectorStoreObservationContext.Builder} instance.
131+
* @param operationName the operation name
132+
* @return the observation context builder
133+
*/
134+
public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);
135+
136+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,7 @@ public void afterPropertiesSet() throws Exception {
387387
}
388388

389389
@Override
390-
public Builder createObservationContextBuilder(String operationName) {
390+
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
391391

392392
return VectorStoreObservationContext.builder(VectorStoreProvider.AZURE.value(), operationName)
393393
.withCollectionName(this.indexName)

‎vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/CassandraVectorStore.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -383,7 +383,7 @@ private String getDocumentId(Row row) {
383383
}
384384

385385
@Override
386-
public Builder createObservationContextBuilder(String operationName) {
386+
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
387387
return VectorStoreObservationContext.builder(VectorStoreProvider.CASSANDRA.value(), operationName)
388388
.withCollectionName(this.conf.schema.table())
389389
.withDimensions(this.embeddingModel.dimensions())

0 commit comments

Comments
 (0)