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
This commit is contained in:
Christian Tzolov
committed by
Mark Pollack
Mark Pollack
parent
72369d515f
commit
bc3f9acb86
@@ -88,6 +88,12 @@
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>io.micrometer</groupId>
|
||||
<artifactId>micrometer-observation-test</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
</project>
|
||||
|
||||
@@ -28,7 +28,12 @@ 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.filter.FilterExpressionConverter;
|
||||
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
|
||||
import org.springframework.beans.factory.InitializingBean;
|
||||
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
@@ -42,6 +47,8 @@ import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.pgvector.PGvector;
|
||||
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
|
||||
/**
|
||||
* Uses the "vector_store" table to store the Spring AI vector data. The table and the
|
||||
* vector index will be auto-created if not available.
|
||||
@@ -50,7 +57,7 @@ import com.pgvector.PGvector;
|
||||
* @author Josh Long
|
||||
* @author Muthukumaran Navaneethakrishnan
|
||||
*/
|
||||
public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class);
|
||||
|
||||
@@ -58,13 +65,13 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
|
||||
public static final int INVALID_EMBEDDING_DIMENSION = -1;
|
||||
|
||||
public final static String DEFAULT_TABLE_NAME = "vector_store";
|
||||
public static final String DEFAULT_TABLE_NAME = "vector_store";
|
||||
|
||||
public final static String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index";
|
||||
public static final String DEFAULT_VECTOR_INDEX_NAME = "spring_ai_vector_index";
|
||||
|
||||
public final static String DEFAULT_SCHEMA_NAME = "public";
|
||||
public static final String DEFAULT_SCHEMA_NAME = "public";
|
||||
|
||||
public final static boolean DEFAULT_SCHEMA_VALIDATION = false;
|
||||
public static final boolean DEFAULT_SCHEMA_VALIDATION = false;
|
||||
|
||||
public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
|
||||
|
||||
@@ -124,10 +131,23 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType,
|
||||
boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema) {
|
||||
|
||||
this(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate, embeddingModel, dimensions,
|
||||
distanceType, removeExistingVectorStoreTable, createIndexMethod, initializeSchema,
|
||||
ObservationRegistry.NOOP, null);
|
||||
}
|
||||
|
||||
private PgVectorStore(String schemaName, String vectorTableName, boolean vectorTableValidationsEnabled,
|
||||
JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel, int dimensions, PgDistanceType distanceType,
|
||||
boolean removeExistingVectorStoreTable, PgIndexType createIndexMethod, boolean initializeSchema,
|
||||
ObservationRegistry observationRegistry,
|
||||
VectorStoreObservationConvention customSearchObservationConvention) {
|
||||
|
||||
super(observationRegistry, customSearchObservationConvention);
|
||||
|
||||
this.vectorTableName = (null == vectorTableName || vectorTableName.isEmpty()) ? DEFAULT_TABLE_NAME
|
||||
: vectorTableName.trim();
|
||||
logger.info("Using the vector table name: {}",
|
||||
this.vectorTableName + " is empty" + (null == vectorTableName || vectorTableName.isEmpty()));
|
||||
logger.info("Using the vector table name: {}. Is empty: {}", this.vectorTableName,
|
||||
(vectorTableName == null || vectorTableName.isEmpty()));
|
||||
|
||||
this.vectorIndexName = this.vectorTableName.equals(DEFAULT_TABLE_NAME) ? DEFAULT_VECTOR_INDEX_NAME
|
||||
: this.vectorTableName + "_index";
|
||||
@@ -150,7 +170,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
}
|
||||
|
||||
@Override
|
||||
public void add(List<Document> documents) {
|
||||
public void doAdd(List<Document> documents) {
|
||||
|
||||
int size = documents.size();
|
||||
|
||||
@@ -195,17 +215,8 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
}
|
||||
}
|
||||
|
||||
private float[] toFloatArray(List<Float> embedding) {
|
||||
float[] embeddingFloat = new float[embedding.size()];
|
||||
int i = 0;
|
||||
for (Float d : embedding) {
|
||||
embeddingFloat[i++] = d.floatValue();
|
||||
}
|
||||
return embeddingFloat;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<Boolean> delete(List<String> idList) {
|
||||
public Optional<Boolean> doDelete(List<String> idList) {
|
||||
int updateCount = 0;
|
||||
for (String id : idList) {
|
||||
int count = jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?",
|
||||
@@ -217,7 +228,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> similaritySearch(SearchRequest request) {
|
||||
public List<Document> doSimilaritySearch(SearchRequest request) {
|
||||
|
||||
String nativeFilterExpression = (request.getFilterExpression() != null)
|
||||
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
|
||||
@@ -276,7 +287,7 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
}
|
||||
|
||||
if (!this.initializeSchema) {
|
||||
logger.debug("Skipping the schema initialization for the table: " + this.getFullyQualifiedTableName());
|
||||
logger.debug("Skipping the schema initialization for the table: {}", this.getFullyQualifiedTableName());
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -484,6 +495,11 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
|
||||
private boolean initializeSchema;
|
||||
|
||||
private ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
|
||||
|
||||
@Nullable
|
||||
private VectorStoreObservationConvention searchObservationConvention;
|
||||
|
||||
// Builder constructor with mandatory parameters
|
||||
public Builder(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
|
||||
if (jdbcTemplate == null || embeddingModel == null) {
|
||||
@@ -533,12 +549,46 @@ public class PgVectorStore implements VectorStore, InitializingBean {
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withObservationRegistry(ObservationRegistry observationRegistry) {
|
||||
this.observationRegistry = observationRegistry;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder withSearchObservationConvention(
|
||||
VectorStoreObservationConvention customSearchObservationConvention) {
|
||||
this.searchObservationConvention = customSearchObservationConvention;
|
||||
return this;
|
||||
}
|
||||
|
||||
public PgVectorStore build() {
|
||||
return new PgVectorStore(schemaName, vectorTableName, vectorTableValidationsEnabled, jdbcTemplate,
|
||||
embeddingModel, dimensions, distanceType, removeExistingVectorStoreTable, indexType,
|
||||
initializeSchema);
|
||||
initializeSchema, observationRegistry, searchObservationConvention);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
|
||||
|
||||
return VectorStoreObservationContext.builder(VectorStoreProvider.PG_VECTOR.value(), operationName)
|
||||
.withDimensions(this.embeddingDimensions())
|
||||
.withCollectionName(this.vectorTableName)
|
||||
.withNamespace(this.schemaName)
|
||||
.withSimilarityMetric(getSimilarityMetric())
|
||||
.withIndexName(this.createIndexMethod.name());
|
||||
}
|
||||
|
||||
private static Map<PgDistanceType, VectorStoreSimilarityMetric> SIMILARITY_TYPE_MAPPING = Map.of(
|
||||
PgDistanceType.COSINE_DISTANCE, VectorStoreSimilarityMetric.COSINE, PgDistanceType.EUCLIDEAN_DISTANCE,
|
||||
VectorStoreSimilarityMetric.EUCLIDEAN, PgDistanceType.NEGATIVE_INNER_PRODUCT,
|
||||
VectorStoreSimilarityMetric.DOT);
|
||||
|
||||
private String getSimilarityMetric() {
|
||||
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.getDistanceType())) {
|
||||
return this.getDistanceType().name();
|
||||
}
|
||||
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
/*
|
||||
* Copyright 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;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import javax.sql.DataSource;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
||||
import org.springframework.ai.openai.api.OpenAiApi;
|
||||
import org.springframework.ai.vectorstore.PgVectorStore.PgIndexType;
|
||||
import org.springframework.ai.vectorstore.observation.DefaultVectorStoreObservationConvention;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.HighCardinalityKeyNames;
|
||||
import org.springframework.ai.vectorstore.observation.VectorStoreObservationDocumentation.LowCardinalityKeyNames;
|
||||
import org.springframework.boot.SpringBootConfiguration;
|
||||
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
|
||||
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
|
||||
import org.springframework.boot.context.properties.ConfigurationProperties;
|
||||
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Primary;
|
||||
import org.springframework.core.io.DefaultResourceLoader;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
import org.testcontainers.containers.PostgreSQLContainer;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
|
||||
import com.zaxxer.hikari.HikariDataSource;
|
||||
|
||||
import io.micrometer.observation.ObservationRegistry;
|
||||
import io.micrometer.observation.tck.TestObservationRegistry;
|
||||
import io.micrometer.observation.tck.TestObservationRegistryAssert;
|
||||
|
||||
/**
|
||||
* Integration tests for observation instruAbstractObservationVectorStorementation in
|
||||
* {@link OpenAiChatModel}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
||||
@Testcontainers
|
||||
public class PgVectorObservationIT {
|
||||
|
||||
@Container
|
||||
@SuppressWarnings("resource")
|
||||
static PostgreSQLContainer<?> postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16")
|
||||
.withUsername("postgres")
|
||||
.withPassword("postgres");
|
||||
|
||||
List<Document> documents = List.of(
|
||||
new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),
|
||||
new Document(getText("classpath:/test/data/time.shelter.txt")),
|
||||
new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("meta2", "meta2")));
|
||||
|
||||
public static String getText(String uri) {
|
||||
var resource = new DefaultResourceLoader().getResource(uri);
|
||||
try {
|
||||
return resource.getContentAsString(StandardCharsets.UTF_8);
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
|
||||
.withUserConfiguration(Config.class)
|
||||
.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE",
|
||||
|
||||
// JdbcTemplate configuration
|
||||
String.format("app.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(),
|
||||
postgresContainer.getMappedPort(5432), "postgres"),
|
||||
"app.datasource.username=postgres", "app.datasource.password=postgres",
|
||||
"app.datasource.type=com.zaxxer.hikari.HikariDataSource");
|
||||
|
||||
@Test
|
||||
void observationVectorStoreAddAndQueryOperations() {
|
||||
|
||||
contextRunner.run(context -> {
|
||||
|
||||
VectorStore vectorStore = context.getBean(VectorStore.class);
|
||||
|
||||
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
|
||||
|
||||
vectorStore.add(documents);
|
||||
|
||||
TestObservationRegistryAssert.assertThat(observationRegistry)
|
||||
.doesNotHaveAnyRemainingCurrentObservation()
|
||||
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
|
||||
.that()
|
||||
.hasContextualNameEqualTo("vector_store pg_vector add")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "add")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), "pg_vector")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), "vector_store")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.QUERY.asString(), "none")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DIMENSIONS.asString(), "1536")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.COLLECTION_NAME.asString(), "vector_store")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.NAMESPACE.asString(), "public")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.FIELD_NAME.asString(), "none")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), "cosine")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.TOP_K.asString(), "none")
|
||||
|
||||
.hasBeenStarted()
|
||||
.hasBeenStopped();
|
||||
|
||||
observationRegistry.clear();
|
||||
|
||||
List<Document> results = vectorStore
|
||||
.similaritySearch(SearchRequest.query("What is Great Depression").withTopK(1));
|
||||
|
||||
assertThat(results).isNotEmpty();
|
||||
|
||||
TestObservationRegistryAssert.assertThat(observationRegistry)
|
||||
.doesNotHaveAnyRemainingCurrentObservation()
|
||||
.hasObservationWithNameEqualTo(DefaultVectorStoreObservationConvention.DEFAULT_NAME)
|
||||
.that()
|
||||
.hasContextualNameEqualTo("vector_store pg_vector query")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_OPERATION_NAME.asString(), "query")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.DB_SYSTEM.asString(), "pg_vector")
|
||||
.hasLowCardinalityKeyValue(LowCardinalityKeyNames.SPRING_AI_KIND.asString(), "vector_store")
|
||||
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.QUERY.asString(), "What is Great Depression")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.DIMENSIONS.asString(), "1536")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.COLLECTION_NAME.asString(), "vector_store")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.NAMESPACE.asString(), "public")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.FIELD_NAME.asString(), "none")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.SIMILARITY_METRIC.asString(), "cosine")
|
||||
.hasHighCardinalityKeyValue(HighCardinalityKeyNames.TOP_K.asString(), "1")
|
||||
|
||||
.hasBeenStarted()
|
||||
.hasBeenStopped();
|
||||
|
||||
});
|
||||
}
|
||||
|
||||
@SpringBootConfiguration
|
||||
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
|
||||
static class Config {
|
||||
|
||||
@Bean
|
||||
public TestObservationRegistry observationRegistry() {
|
||||
return TestObservationRegistry.create();
|
||||
}
|
||||
|
||||
@Bean
|
||||
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel,
|
||||
ObservationRegistry observationRegistry) {
|
||||
return new PgVectorStore.Builder(jdbcTemplate, embeddingModel)
|
||||
.withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)
|
||||
.withIndexType(PgIndexType.HNSW)
|
||||
.withObservationRegistry(observationRegistry)
|
||||
.withInitializeSchema(true)
|
||||
.build();
|
||||
}
|
||||
|
||||
@Bean
|
||||
public JdbcTemplate myJdbcTemplate(DataSource dataSource) {
|
||||
return new JdbcTemplate(dataSource);
|
||||
}
|
||||
|
||||
@Bean
|
||||
@Primary
|
||||
@ConfigurationProperties("app.datasource")
|
||||
public DataSourceProperties dataSourceProperties() {
|
||||
return new DataSourceProperties();
|
||||
}
|
||||
|
||||
@Bean
|
||||
public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) {
|
||||
return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build();
|
||||
}
|
||||
|
||||
@Bean
|
||||
public EmbeddingModel embeddingModel() {
|
||||
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
Reference in New Issue
Block a user