Skip to content

Commit

Permalink
Add observability support to VectorStore
Browse files Browse the repository at this point in the history
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
  • Loading branch information
tzolov authored and markpollack committed Aug 16, 2024
1 parent 72369d5 commit bc3f9ac
Show file tree
Hide file tree
Showing 23 changed files with 1,620 additions and 38 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,11 @@ public enum AiObservationAttributes {
/**
* The full response received from the model.
*/
COMPLETION("gen_ai.completion");
COMPLETION("gen_ai.completion"),
/**
* The name of the operation or command being executed.
*/
DB_OPERATION_NAME("db.operation.name"),;

private final String value;

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.observation.conventions;

/**
* @author Christian Tzolov
* @since 1.0.0
*/
public enum VectorStoreProvider {

// @formatter:off
PG_VECTOR("pg_vector"),
SIMPLE_VECTOR_STORE("simple_vector_store");

// @formatter:on
private final String value;

VectorStoreProvider(String value) {
this.value = value;
}

public String value() {
return this.value;
}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.observation.conventions;

/**
* @author Christian Tzolov
* @since 1.0.0
*/
public enum VectorStoreSimilarityMetric {

// @formatter:off

/**
* The cosine metric.
*/
COSINE("cosine"),
/**
* The euclidean distance metric.
*/
EUCLIDEAN("euclidean"),
/**
* The manhattan distance metric.
*/
MANHATTAN("manhattan"),
/**
* The dot product metric.
*/
DOT("dot");

// @formatter:on
private final String value;

VectorStoreSimilarityMetric(String value) {
this.value = value;
}

public String value() {
return this.value;
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -15,23 +15,15 @@
*/
package org.springframework.ai.vectorstore;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.core.io.Resource;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.FileAlreadyExistsException;
import java.nio.file.Files;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
Expand All @@ -40,6 +32,21 @@
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.observation.conventions.VectorStoreProvider;
import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetric;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.core.io.Resource;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;

/**
* SimpleVectorStore is a simple implementation of the VectorStore interface.
*
Expand All @@ -55,7 +62,7 @@
* @author Mark Pollack
* @author Christian Tzolov
*/
public class SimpleVectorStore implements VectorStore {
public class SimpleVectorStore extends AbstractObservationVectorStore {

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

Expand All @@ -69,7 +76,7 @@ public SimpleVectorStore(EmbeddingModel embeddingModel) {
}

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

@Override
public Optional<Boolean> delete(List<String> idList) {
public Optional<Boolean> doDelete(List<String> idList) {
for (String id : idList) {
this.store.remove(id);
}
return Optional.of(true);
}

@Override
public List<Document> similaritySearch(SearchRequest request) {
public List<Document> doSimilaritySearch(SearchRequest request) {
if (request.getFilterExpression() != null) {
throw new UnsupportedOperationException(
"The [" + this.getClass() + "] doesn't support metadata filtering!");
Expand All @@ -114,7 +121,15 @@ public void save(File file) {
try {
if (!file.exists()) {
logger.info("Creating new vector store file: {}", file);
file.createNewFile();
try {
Files.createFile(file.toPath());
}
catch (FileAlreadyExistsException e) {
throw new RuntimeException("File already exists: " + file, e);
}
catch (IOException e) {
throw new RuntimeException("Failed to create new file: " + file + ". Reason: " + e.getMessage(), e);
}
}
else {
logger.info("Overwriting existing vector store file: {}", file);
Expand Down Expand Up @@ -247,4 +262,13 @@ public static float norm(float[] vector) {

}

@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {

return VectorStoreObservationContext.builder(VectorStoreProvider.SIMPLE_VECTOR_STORE.value(), operationName)
.withDimensions(this.embeddingModel.dimensions())
.withCollectionName("in-memory-map")
.withSimilarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,10 @@
*/
public interface VectorStore extends DocumentWriter {

default String getName() {
return this.getClass().getSimpleName();
}

/**
* Adds list of {@link Document}s to the vector store.
* @param documents the list of documents to store. Throws an exception if the
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore.observation;

import java.util.List;
import java.util.Optional;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.lang.Nullable;

import io.micrometer.observation.ObservationRegistry;

/**
* @author Christian Tzolov
* @since 1.0.0
*/
public abstract class AbstractObservationVectorStore implements VectorStore {

private static final VectorStoreObservationConvention DEFAULT_OBSERVATION_CONVENTION = new DefaultVectorStoreObservationConvention();

private final ObservationRegistry observationRegistry;

@Nullable
private final VectorStoreObservationConvention customObservationConvention;

public AbstractObservationVectorStore() {
this(ObservationRegistry.NOOP, null);
}

public AbstractObservationVectorStore(ObservationRegistry observationRegistry) {
this(observationRegistry, null);
}

public AbstractObservationVectorStore(ObservationRegistry observationRegistry,
VectorStoreObservationConvention customSearchObservationConvention) {
this.observationRegistry = observationRegistry;
this.customObservationConvention = customSearchObservationConvention;
}

@Override
public void add(List<Document> documents) {

VectorStoreObservationContext observationContext = this
.createObservationContextBuilder(VectorStoreObservationContext.Operation.ADD.value())
.build();

VectorStoreObservationDocumentation.AI_VECTOR_STORE
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
observationRegistry)
.observe(() -> this.doAdd(documents));
}

@Override
public Optional<Boolean> delete(List<String> deleteDocIds) {

VectorStoreObservationContext observationContext = this
.createObservationContextBuilder(VectorStoreObservationContext.Operation.DELETE.value())
.build();

return VectorStoreObservationDocumentation.AI_VECTOR_STORE
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry)
.observe(() -> this.doDelete(deleteDocIds));
}

@Override
public List<Document> similaritySearch(SearchRequest request) {

VectorStoreObservationContext searchObservationContext = this
.createObservationContextBuilder(VectorStoreObservationContext.Operation.QUERY.value())
.withQueryRequest(request)
.build();

return VectorStoreObservationDocumentation.AI_VECTOR_STORE
.observation(this.customObservationConvention, DEFAULT_OBSERVATION_CONVENTION,
() -> searchObservationContext, this.observationRegistry)
.observe(() -> {
var documents = this.doSimilaritySearch(request);
searchObservationContext.setQueryResponse(documents);
return documents;
});
}

public abstract void doAdd(List<Document> documents);

public abstract Optional<Boolean> doDelete(List<String> idList);

public abstract List<Document> doSimilaritySearch(SearchRequest request);

public abstract VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName);

}
Loading

0 comments on commit bc3f9ac

Please sign in to comment.