Skip to content

Commit bc3f9ac

Browse files
tzolovmarkpollack
authored andcommitted
Add observability support to VectorStore
Implementation: - Introduce AbstractObservationVectorStore with instrumentation for add, delete, and similaritySearch methods - Create VectorStoreObservationContext to capture operation details - Implement DefaultVectorStoreObservationConvention for naming and tagging - Add VectorStoreObservationDocumentation for defining observation keys - Create VectorStoreObservationAutoConfiguration for auto-configuring observations - Add VectorStoreObservationProperties to control optional observation content filters - Update VectorStore interface with getName() method - Modify PgVectorStore and SimpleVectorStore to extend AbstractObservationVectorStore - Add vector_store Spring AI kind Filters: - Implement VectorStoreQueryResponseObservationFilter - Add VectorStoreDeleteRequestContentObservationFilter and VectorStoreAddRequestContentObservationFilter Enhancements: - Update PgVectorStoreAutoConfiguration to support observations - Add observation support to PgVectorStore's Builder - Add VectorStoreObservationContext.Operation enum with ADD, DELETE, and QUERY options Tests: - Add tests for VectorStore context, convention, and filters - Add VectorStoreObservationAutoConfiguration tests - Add PgVectorObservationIT Resolves #1205
1 parent 72369d5 commit bc3f9ac

File tree

23 files changed

+1620
-38
lines changed

23 files changed

+1620
-38
lines changed

spring-ai-core/src/main/java/org/springframework/ai/observation/conventions/AiObservationAttributes.java

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,11 @@ public enum AiObservationAttributes {
132132
/**
133133
* The full response received from the model.
134134
*/
135-
COMPLETION("gen_ai.completion");
135+
COMPLETION("gen_ai.completion"),
136+
/**
137+
* The name of the operation or command being executed.
138+
*/
139+
DB_OPERATION_NAME("db.operation.name"),;
136140

137141
private final String value;
138142

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2024 - 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+
package org.springframework.ai.observation.conventions;
17+
18+
/**
19+
* @author Christian Tzolov
20+
* @since 1.0.0
21+
*/
22+
public enum VectorStoreProvider {
23+
24+
// @formatter:off
25+
PG_VECTOR("pg_vector"),
26+
SIMPLE_VECTOR_STORE("simple_vector_store");
27+
28+
// @formatter:on
29+
private final String value;
30+
31+
VectorStoreProvider(String value) {
32+
this.value = value;
33+
}
34+
35+
public String value() {
36+
return this.value;
37+
}
38+
39+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright 2024 - 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+
package org.springframework.ai.observation.conventions;
17+
18+
/**
19+
* @author Christian Tzolov
20+
* @since 1.0.0
21+
*/
22+
public enum VectorStoreSimilarityMetric {
23+
24+
// @formatter:off
25+
26+
/**
27+
* The cosine metric.
28+
*/
29+
COSINE("cosine"),
30+
/**
31+
* The euclidean distance metric.
32+
*/
33+
EUCLIDEAN("euclidean"),
34+
/**
35+
* The manhattan distance metric.
36+
*/
37+
MANHATTAN("manhattan"),
38+
/**
39+
* The dot product metric.
40+
*/
41+
DOT("dot");
42+
43+
// @formatter:on
44+
private final String value;
45+
46+
VectorStoreSimilarityMetric(String value) {
47+
this.value = value;
48+
}
49+
50+
public String value() {
51+
return this.value;
52+
}
53+
54+
}

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

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

18-
import com.fasterxml.jackson.core.JsonProcessingException;
19-
import com.fasterxml.jackson.core.type.TypeReference;
20-
import com.fasterxml.jackson.databind.ObjectMapper;
21-
import com.fasterxml.jackson.databind.ObjectWriter;
22-
import org.slf4j.Logger;
23-
import org.slf4j.LoggerFactory;
24-
import org.springframework.ai.document.Document;
25-
import org.springframework.ai.embedding.EmbeddingModel;
26-
import org.springframework.core.io.Resource;
27-
2818
import java.io.File;
2919
import java.io.FileOutputStream;
3020
import java.io.IOException;
3121
import java.io.OutputStream;
3222
import java.io.OutputStreamWriter;
3323
import java.io.Writer;
3424
import java.nio.charset.StandardCharsets;
25+
import java.nio.file.FileAlreadyExistsException;
26+
import java.nio.file.Files;
3527
import java.util.Comparator;
3628
import java.util.HashMap;
3729
import java.util.List;
@@ -40,6 +32,21 @@
4032
import java.util.Optional;
4133
import java.util.concurrent.ConcurrentHashMap;
4234

35+
import org.slf4j.Logger;
36+
import org.slf4j.LoggerFactory;
37+
import org.springframework.ai.document.Document;
38+
import org.springframework.ai.embedding.EmbeddingModel;
39+
import org.springframework.ai.observation.conventions.VectorStoreProvider;
40+
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
41+
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
42+
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
43+
import org.springframework.core.io.Resource;
44+
45+
import com.fasterxml.jackson.core.JsonProcessingException;
46+
import com.fasterxml.jackson.core.type.TypeReference;
47+
import com.fasterxml.jackson.databind.ObjectMapper;
48+
import com.fasterxml.jackson.databind.ObjectWriter;
49+
4350
/**
4451
* SimpleVectorStore is a simple implementation of the VectorStore interface.
4552
*
@@ -55,7 +62,7 @@
5562
* @author Mark Pollack
5663
* @author Christian Tzolov
5764
*/
58-
public class SimpleVectorStore implements VectorStore {
65+
public class SimpleVectorStore extends AbstractObservationVectorStore {
5966

6067
private static final Logger logger = LoggerFactory.getLogger(SimpleVectorStore.class);
6168

@@ -69,7 +76,7 @@ public SimpleVectorStore(EmbeddingModel embeddingModel) {
6976
}
7077

7178
@Override
72-
public void add(List<Document> documents) {
79+
public void doAdd(List<Document> documents) {
7380
for (Document document : documents) {
7481
logger.info("Calling EmbeddingModel for document id = {}", document.getId());
7582
float[] embedding = this.embeddingModel.embed(document);
@@ -79,15 +86,15 @@ public void add(List<Document> documents) {
7986
}
8087

8188
@Override
82-
public Optional<Boolean> delete(List<String> idList) {
89+
public Optional<Boolean> doDelete(List<String> idList) {
8390
for (String id : idList) {
8491
this.store.remove(id);
8592
}
8693
return Optional.of(true);
8794
}
8895

8996
@Override
90-
public List<Document> similaritySearch(SearchRequest request) {
97+
public List<Document> doSimilaritySearch(SearchRequest request) {
9198
if (request.getFilterExpression() != null) {
9299
throw new UnsupportedOperationException(
93100
"The [" + this.getClass() + "] doesn't support metadata filtering!");
@@ -114,7 +121,15 @@ public void save(File file) {
114121
try {
115122
if (!file.exists()) {
116123
logger.info("Creating new vector store file: {}", file);
117-
file.createNewFile();
124+
try {
125+
Files.createFile(file.toPath());
126+
}
127+
catch (FileAlreadyExistsException e) {
128+
throw new RuntimeException("File already exists: " + file, e);
129+
}
130+
catch (IOException e) {
131+
throw new RuntimeException("Failed to create new file: " + file + ". Reason: " + e.getMessage(), e);
132+
}
118133
}
119134
else {
120135
logger.info("Overwriting existing vector store file: {}", file);
@@ -247,4 +262,13 @@ public static float norm(float[] vector) {
247262

248263
}
249264

265+
@Override
266+
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
267+
268+
return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE_VECTOR_STORE.value(), operationName)
269+
.withDimensions(this.embeddingModel.dimensions())
270+
.withCollectionName("in-memory-map")
271+
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
272+
}
273+
250274
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@
3131
*/
3232
public interface VectorStore extends DocumentWriter {
3333

34+
default String getName() {
35+
return this.getClass().getSimpleName();
36+
}
37+
3438
/**
3539
* Adds list of {@link Document}s to the vector store.
3640
* @param documents the list of documents to store. Throws an exception if the
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
/*
2+
* Copyright 2024 - 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+
package org.springframework.ai.vectorstore.observation;
17+
18+
import java.util.List;
19+
import java.util.Optional;
20+
21+
import org.springframework.ai.document.Document;
22+
import org.springframework.ai.vectorstore.SearchRequest;
23+
import org.springframework.ai.vectorstore.VectorStore;
24+
import org.springframework.lang.Nullable;
25+
26+
import io.micrometer.observation.ObservationRegistry;
27+
28+
/**
29+
* @author Christian Tzolov
30+
* @since 1.0.0
31+
*/
32+
public abstract class AbstractObservationVectorStore implements VectorStore {
33+
34+
private static final VectorStoreObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultVectorStoreObservationConvention();
35+
36+
private final ObservationRegistry observationRegistry;
37+
38+
@Nullable
39+
private final VectorStoreObservationConvention customObservationConvention;
40+
41+
public AbstractObservationVectorStore() {
42+
this(ObservationRegistry.NOOP, null);
43+
}
44+
45+
public AbstractObservationVectorStore(ObservationRegistry observationRegistry) {
46+
this(observationRegistry, null);
47+
}
48+
49+
public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
50+
VectorStoreObservationConvention customSearchObservationConvention) {
51+
this.observationRegistry = observationRegistry;
52+
this.customObservationConvention = customSearchObservationConvention;
53+
}
54+
55+
@Override
56+
public void add(List<Document> documents) {
57+
58+
VectorStoreObservationContext observationContext = this
59+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.ADD.value())
60+
.build();
61+
62+
VectorStoreObservationDocumentation.AI_VECTOR_STORE
63+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
64+
observationRegistry)
65+
.observe(() -> this.doAdd(documents));
66+
}
67+
68+
@Override
69+
public Optional<Boolean> delete(List<String> deleteDocIds) {
70+
71+
VectorStoreObservationContext observationContext = this
72+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value())
73+
.build();
74+
75+
return VectorStoreObservationDocumentation.AI_VECTOR_STORE
76+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
77+
this.observationRegistry)
78+
.observe(() -> this.doDelete(deleteDocIds));
79+
}
80+
81+
@Override
82+
public List<Document> similaritySearch(SearchRequest request) {
83+
84+
VectorStoreObservationContext searchObservationContext = this
85+
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
86+
.withQueryRequest(request)
87+
.build();
88+
89+
return VectorStoreObservationDocumentation.AI_VECTOR_STORE
90+
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION,
91+
() -> searchObservationContext, this.observationRegistry)
92+
.observe(() -> {
93+
var documents = this.doSimilaritySearch(request);
94+
searchObservationContext.setQueryResponse(documents);
95+
return documents;
96+
});
97+
}
98+
99+
public abstract void doAdd(List<Document> documents);
100+
101+
public abstract Optional<Boolean> doDelete(List<String> idList);
102+
103+
public abstract List<Document> doSimilaritySearch(SearchRequest request);
104+
105+
public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);
106+
107+
}

0 commit comments

Comments
 (0)