Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Introduce checkstyle plugin

Browse Source 
- Based on https://github.com/spring-io/spring-javaformat
- In this iteration, checkstyles are only enabled for spring-ai-core
This commit is contained in:
Soby Chacko
2024-10-24 10:39:48 -04:00
committed by Mark Pollack
parent 33a72417e1
commit 8e758dbd00
1412 changed files with 26997 additions and 21963 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,14 +13,16 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter;
import java.util.List;
/**
* Converts {@link Expression} into PgVector metadata filter expression format.

10
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorSchemaValidator.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.ArrayList;
@@ -21,6 +22,7 @@ import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.JdbcTemplate;
@@ -64,7 +66,7 @@ public class PgVectorSchemaValidator {
String sql = "SELECT 1 FROM information_schema.tables WHERE table_schema = ? AND table_name = ?";
try {
// Query for a single integer value, if it exists, table exists
jdbcTemplate.queryForObject(sql, Integer.class, schemaName, tableName);
this.jdbcTemplate.queryForObject(sql, Integer.class, schemaName, tableName);
return true;
}
catch (DataAccessException e) {
@@ -100,7 +102,7 @@ public class PgVectorSchemaValidator {
// Include the schema name in the query to target the correct table
String query = "SELECT column_name, data_type FROM information_schema.columns "
+ "WHERE table_schema = ? AND table_name = ?";
List<Map<String, Object>> columns = jdbcTemplate.queryForList(query,
List<Map<String, Object>> columns = this.jdbcTemplate.queryForList(query,
new Object[] { schemaName, tableName });
if (columns.isEmpty()) {

86
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorStore.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,8 +13,18 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
@@ -23,6 +33,7 @@ import io.micrometer.observation.ObservationRegistry;
import org.postgresql.util.PGobject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -44,15 +55,6 @@ import org.springframework.jdbc.core.StatementCreatorUtils;
import org.springframework.lang.Nullable;
import org.springframework.util.StringUtils;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
/**
* 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.
@@ -67,8 +69,6 @@ import java.util.UUID;
*/
public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class);
public static final int OPENAI_EMBEDDING_DIMENSION_SIZE = 1536;
public static final int INVALID_EMBEDDING_DIMENSION = -1;
@@ -81,10 +81,17 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public static final boolean DEFAULT_SCHEMA_VALIDATION = false;
public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
public static final int MAX_DOCUMENT_BATCH_SIZE = 10_000;
private static final Logger logger = LoggerFactory.getLogger(PgVectorStore.class);
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);
public final FilterExpressionConverter filterExpressionConverter = new PgVectorFilterExpressionConverter();
private final String vectorTableName;
private final String vectorIndexName;
@@ -183,7 +190,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
}
public PgDistanceType getDistanceType() {
return distanceType;
return this.distanceType;
}
@Override
@@ -208,6 +215,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
+ "UPDATE SET content = ? , metadata = ?::jsonb , embedding = ? ";
this.jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
@@ -247,7 +255,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
public Optional<Boolean> doDelete(List<String> idList) {
int updateCount = 0;
for (String id : idList) {
int count = jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?",
int count = this.jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?",
UUID.fromString(id));
updateCount = updateCount + count;
}
@@ -281,6 +289,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return this.jdbcTemplate.query(
"SELECT embedding " + this.comparisonOperator() + " ? AS distance FROM " + getFullyQualifiedTableName(),
new RowMapper<Double>() {
@Override
@Nullable
public Double mapRow(ResultSet rs, int rowNum) throws SQLException {
@@ -383,6 +392,23 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return OPENAI_EMBEDDING_DIMENSION_SIZE;
}
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.PG_VECTOR.value(), operationName)
.withCollectionName(this.vectorTableName)
.withDimensions(this.embeddingDimensions())
.withNamespace(this.schemaName)
.withSimilarityMetric(getSimilarityMetric());
}
private String getSimilarityMetric() {
if (!SIMILARITY_TYPE_MAPPING.containsKey(this.getDistanceType())) {
return this.getDistanceType().name();
}
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
/**
* By default, pgvector performs exact nearest neighbor search, which provides perfect
* recall. You can add an index to use approximate nearest neighbor search, which
@@ -492,7 +518,7 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
String source = pgObject.getValue();
try {
return (Map<String, Object>) objectMapper.readValue(source, Map.class);
return (Map<String, Object>) this.objectMapper.readValue(source, Map.class);
}
catch (JsonProcessingException e) {
throw new RuntimeException(e);
@@ -611,26 +637,4 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
}
@Override
public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
return VectorStoreObservationContext.builder(VectorStoreProvider.PG_VECTOR.value(), operationName)
.withCollectionName(this.vectorTableName)
.withDimensions(this.embeddingDimensions())
.withNamespace(this.schemaName)
.withSimilarityMetric(getSimilarityMetric());
}
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();
}
}
}

21
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorEmbeddingDimensionsTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.junit.jupiter.api.Test;
@@ -46,32 +47,32 @@ public class PgVectorEmbeddingDimensionsTests {
final int explicitDimensions = 696;
var dim = new PgVectorStore(jdbcTemplate, embeddingModel, explicitDimensions).embeddingDimensions();
var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel, explicitDimensions).embeddingDimensions();
assertThat(dim).isEqualTo(explicitDimensions);
verify(embeddingModel, never()).dimensions();
verify(this.embeddingModel, never()).dimensions();
}
@Test
public void embeddingModelDimensions() {
when(embeddingModel.dimensions()).thenReturn(969);
when(this.embeddingModel.dimensions()).thenReturn(969);
var dim = new PgVectorStore(jdbcTemplate, embeddingModel).embeddingDimensions();
var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel).embeddingDimensions();
assertThat(dim).isEqualTo(969);
verify(embeddingModel, only()).dimensions();
verify(this.embeddingModel, only()).dimensions();
}
@Test
public void fallBackToDefaultDimensions() {
when(embeddingModel.dimensions()).thenThrow(new RuntimeException());
when(this.embeddingModel.dimensions()).thenThrow(new RuntimeException());
var dim = new PgVectorStore(jdbcTemplate, embeddingModel).embeddingDimensions();
var dim = new PgVectorStore(this.jdbcTemplate, this.embeddingModel).embeddingDimensions();
assertThat(dim).isEqualTo(PgVectorStore.OPENAI_EMBEDDING_DIMENSION_SIZE);
verify(embeddingModel, only()).dimensions();
verify(this.embeddingModel, only()).dimensions();
}
}

33
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,13 +13,20 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vectorstore.filter.Filter.Expression;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.AND;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.GTE;
@@ -28,10 +35,6 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.LT
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NE;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.NIN;
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.OR;
import org.springframework.ai.vectorstore.filter.Filter.Group;
import org.springframework.ai.vectorstore.filter.Filter.Key;
import org.springframework.ai.vectorstore.filter.Filter.Value;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
/**
* @author Muthukumaran Navaneethakrishnan
@@ -44,14 +47,14 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void testEQ() {
// country == "BG"
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
String vectorExpr = this.converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(vectorExpr).isEqualTo("$.country == \"BG\"");
}
@Test
public void tesEqAndGte() {
// genre == "drama" AND year >= 2020
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(vectorExpr).isEqualTo("$.genre == \"drama\" && $.year >= 2020");
@@ -60,7 +63,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void tesIn() {
// genre in ["comedy", "documentary", "drama"]
String vectorExpr = converter.convertExpression(
String vectorExpr = this.converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(vectorExpr)
.isEqualTo("($.genre == \"comedy\" || $.genre == \"documentary\" || $.genre == \"drama\")");
@@ -69,7 +72,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void testNe() {
// year >= 2020 OR country == "BG" AND city != "Sofia"
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
@@ -79,7 +82,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void testGroup() {
// (year >= 2020 OR country == "BG") AND city NIN ["Sofia", "Plovdiv"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv")))));
@@ -90,7 +93,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void tesBoolean() {
// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String vectorExpr = converter.convertExpression(new Expression(AND,
String vectorExpr = this.converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));
@@ -102,7 +105,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void testDecimal() {
// temperature >= -15.6 && temperature <= +20.13
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));
@@ -111,7 +114,7 @@ public class PgVectorFilterExpressionConverterTests {
@Test
public void testComplexIdentifiers() {
String vectorExpr = converter
String vectorExpr = this.converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(vectorExpr).isEqualTo("$.\"country 1 2 3\" == \"BG\"");
}

5
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorImage.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,6 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import org.testcontainers.utility.DockerImageName;

48
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreCustomNamesIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,10 +13,19 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Random;
import javax.sql.DataSource;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
@@ -32,12 +41,6 @@ import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.core.JdbcTemplate;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import javax.sql.DataSource;
import java.util.Random;
import static org.assertj.core.api.Assertions.assertThat;
@@ -92,19 +95,20 @@ public class PgVectorStoreCustomNamesIT {
@Test
public void shouldCreateDefaultTableAndIndexIfNotPresentInConfig() {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.schemaValidation=false").run(context -> {
assertThat(context).hasNotFailed();
assertThat(isTableExists(context, "vector_store")).isTrue();
assertThat(isSchemaExists(context, "public")).isTrue();
dropTableByName(context, "vector_store");
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.schemaValidation=false")
.run(context -> {
assertThat(context).hasNotFailed();
assertThat(isTableExists(context, "vector_store")).isTrue();
assertThat(isSchemaExists(context, "public")).isTrue();
dropTableByName(context, "vector_store");
});
});
}
@Test
public void shouldCreateTableAndIndexIfNotPresentInDatabase() {
String tableName = "new_vector_table";
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName)
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName)
.run(context -> {
assertThat(isTableExists(context, tableName)).isTrue();
assertThat(isIndexExists(context, "public", tableName, tableName + "_index")).isTrue();
@@ -118,7 +122,7 @@ public class PgVectorStoreCustomNamesIT {
String tableName = "customvectortable";
contextRunner
this.contextRunner
.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName,
"test.spring.ai.vectorstore.pgvector.schemaValidation=true")
@@ -136,7 +140,7 @@ public class PgVectorStoreCustomNamesIT {
String tableName = "users; DROP TABLE users;";
contextRunner
this.contextRunner
.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName,
"test.spring.ai.vectorstore.pgvector.schemaValidation=true")
@@ -156,7 +160,7 @@ public class PgVectorStoreCustomNamesIT {
String schemaName = "public; DROP TABLE users;";
String tableName = "customvectortable";
contextRunner
this.contextRunner
.withPropertyValues("test.spring.ai.vectorstore.pgvector.vectorTableName=" + tableName,
"test.spring.ai.vectorstore.pgvector.schemaName=" + schemaName,
"test.spring.ai.vectorstore.pgvector.schemaValidation=true")
@@ -189,10 +193,10 @@ public class PgVectorStoreCustomNamesIT {
@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(schemaName)
.withVectorTableName(vectorTableName)
.withVectorTableValidationsEnabled(schemaValidation)
.withDimensions(dimensions)
return new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withSchemaName(this.schemaName)
.withVectorTableName(this.vectorTableName)
.withVectorTableValidationsEnabled(this.schemaValidation)
.withDimensions(this.dimensions)
.withDistanceType(PgVectorStore.PgDistanceType.COSINE_DISTANCE)
.withRemoveExistingVectorStoreTable(true)
.withIndexType(PgIndexType.HNSW)

149
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* 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;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -28,12 +27,17 @@ import java.util.stream.Stream;
import javax.sql.DataSource;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.Assert;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
@@ -53,11 +57,8 @@ import org.springframework.context.annotation.Primary;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.util.CollectionUtils;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import com.zaxxer.hikari.HikariDataSource;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Muthukumaran Navaneethakrishnan
@@ -74,6 +75,16 @@ public class PgVectorStoreIT {
.withUsername("postgres")
.withPassword("postgres");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.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");
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")),
@@ -89,50 +100,11 @@ public class PgVectorStoreIT {
}
}
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.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");
private static void dropTable(ApplicationContext context) {
JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class);
jdbcTemplate.execute("DROP TABLE IF EXISTS vector_store");
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" })
public void addAndSearch(String distanceType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("What is Great Depression").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(2).getId());
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results2).hasSize(0);
dropTable(context);
});
}
static Stream<Arguments> provideFilters() {
return Stream.of(Arguments.of("country in ['BG','NL']", 3), // String Filters In
Arguments.of("year in [2020]", 1), // Numeric Filters In
@@ -141,11 +113,60 @@ public class PgVectorStoreIT {
);
}
private static boolean isSortedByDistance(List<Document> docs) {
List<Float> distances = docs.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
if (CollectionUtils.isEmpty(distances) || distances.size() == 1) {
return true;
}
Iterator<Float> iter = distances.iterator();
Float current, previous = iter.next();
while (iter.hasNext()) {
current = iter.next();
if (previous > current) {
return false;
}
previous = current;
}
return true;
}
@ParameterizedTest(name = "{0} : {displayName} ")
@ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" })
public void addAndSearch(String distanceType) {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(this.documents);
List<Document> results = vectorStore
.similaritySearch(SearchRequest.query("What is Great Depression").withTopK(1));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(2).getId());
assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance");
// Remove all documents from the store
vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore
.similaritySearch(SearchRequest.query("Great Depression").withTopK(1));
assertThat(results2).hasSize(0);
dropTable(context);
});
}
@ParameterizedTest(name = "Filter expression {0} should return {1} records ")
@MethodSource("provideFilters")
public void searchWithInFilter(String expression, Integer expectedRecords) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE")
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE")
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -177,7 +198,7 @@ public class PgVectorStoreIT {
@ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" })
public void searchWithFilters(String distanceType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -251,7 +272,7 @@ public class PgVectorStoreIT {
@ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" })
public void documentUpdate(String distanceType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
@@ -292,12 +313,12 @@ public class PgVectorStoreIT {
// @ValueSource(strings = { "COSINE_DISTANCE" })
public void searchWithThreshold(String distanceType) {
contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + distanceType)
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
List<Document> fullResult = vectorStore
.similaritySearch(SearchRequest.query("Time Shelter").withTopK(5).withSimilarityThresholdAll());
@@ -317,32 +338,12 @@ public class PgVectorStoreIT {
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);
assertThat(resultDoc.getId()).isEqualTo(documents.get(1).getId());
assertThat(resultDoc.getId()).isEqualTo(this.documents.get(1).getId());
dropTable(context);
});
}
private static boolean isSortedByDistance(List<Document> docs) {
List<Float> distances = docs.stream().map(doc -> (Float) doc.getMetadata().get("distance")).toList();
if (CollectionUtils.isEmpty(distances) || distances.size() == 1) {
return true;
}
Iterator<Float> iter = distances.iterator();
Float current, previous = iter.next();
while (iter.hasNext()) {
current = iter.next();
if (previous > current) {
return false;
}
previous = current;
}
return true;
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
@@ -353,7 +354,7 @@ public class PgVectorStoreIT {
@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
return new PgVectorStore(jdbcTemplate, embeddingModel, PgVectorStore.INVALID_EMBEDDING_DIMENSION,
distanceType, true, PgIndexType.HNSW, true);
this.distanceType, true, PgIndexType.HNSW, true);
}
@Bean

48
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreObservationIT.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,9 +13,8 @@
* 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;
package org.springframework.ai.vectorstore;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
@@ -24,8 +23,16 @@ import java.util.Map;
import javax.sql.DataSource;
import com.zaxxer.hikari.HikariDataSource;
import io.micrometer.observation.ObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistry;
import io.micrometer.observation.tck.TestObservationRegistryAssert;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.observation.conventions.SpringAiKind;
@@ -48,15 +55,8 @@ 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;
import static org.assertj.core.api.Assertions.assertThat;
/**
* Integration tests for observation instruAbstractObservationVectorStorementation in
@@ -75,6 +75,16 @@ public class PgVectorStoreObservationIT {
.withUsername("postgres")
.withPassword("postgres");
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");
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")),
@@ -90,26 +100,16 @@ public class PgVectorStoreObservationIT {
}
}
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 -> {
this.contextRunner.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
TestObservationRegistry observationRegistry = context.getBean(TestObservationRegistry.class);
vectorStore.add(documents);
vectorStore.add(this.documents);
TestObservationRegistryAssert.assertThat(observationRegistry)
.doesNotHaveAnyRemainingCurrentObservation()

19
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreTests.java
View File

@@ -1,11 +1,11 @@
/*
* Copyright 2023 - 2024 the original author or authors.
* Copyright 2023-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
* 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,
@@ -13,13 +13,21 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.vectorstore;
import java.util.Collections;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;
import org.mockito.ArgumentCaptor;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
@@ -29,13 +37,6 @@ import static org.mockito.Mockito.only;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import java.util.Collections;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
/**
* @author Muthukumaran Navaneethakrishnan
* @author Soby Chacko

141
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreWithChatMemoryAdvisorIT.java
View File

@@ -5,7 +5,7 @@
* 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
* 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,
@@ -16,12 +16,6 @@
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.util.List;
import java.util.Map;
@@ -32,6 +26,10 @@ import org.mockito.ArgumentCaptor;
import org.mockito.ArgumentMatchers;
import org.mockito.Mockito;
import org.postgresql.ds.PGSimpleDataSource;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.VectorStoreChatMemoryAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
@@ -43,9 +41,12 @@ import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.jdbc.core.JdbcTemplate;
import org.testcontainers.containers.PostgreSQLContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
/**
* @author Fabian Krüger
@@ -55,14 +56,68 @@ import org.testcontainers.junit.jupiter.Testcontainers;
@Testcontainers
class PgVectorStoreWithChatMemoryAdvisorIT {
float[] embed = { 0.003961659F, -0.0073295482F, 0.02663665F };
@Container
@SuppressWarnings("resource")
static PostgreSQLContainer<?> postgresContainer = new PostgreSQLContainer<>(PgVectorImage.DEFAULT_IMAGE)
.withUsername("postgres")
.withPassword("postgres");
float[] embed = { 0.003961659F, -0.0073295482F, 0.02663665F };
private static @NotNull ChatModel chatModelAlwaysReturnsTheSameReply() {
ChatModel chatModel = mock(ChatModel.class);
ArgumentCaptor<Prompt> argumentCaptor = ArgumentCaptor.forClass(Prompt.class);
ChatResponse chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
Why don't scientists trust atoms?
Because they make up everything!
"""))));
when(chatModel.call(argumentCaptor.capture())).thenReturn(chatResponse);
return chatModel;
}
private static void initStore(PgVectorStore store) throws Exception {
store.afterPropertiesSet();
// fill the store
store.add(List.of(new Document("Tell me a good joke", Map.of("conversationId", "default")),
new Document("Tell me a bad joke", Map.of("conversationId", "default", "messageType", "USER"))));
}
private static PgVectorStore createPgVectorStoreUsingTestcontainer(EmbeddingModel embeddingModel) throws Exception {
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
PgVectorStore vectorStore = new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withDimensions(3) // match
// embeddings
.withInitializeSchema(true)
.build();
initStore(vectorStore);
return vectorStore;
}
private static @NotNull JdbcTemplate createJdbcTemplateWithConnectionToTestcontainer() {
PGSimpleDataSource ds = new PGSimpleDataSource();
ds.setUrl("jdbc:postgresql://localhost:" + postgresContainer.getMappedPort(5432) + "/postgres");
ds.setUser(postgresContainer.getUsername());
ds.setPassword(postgresContainer.getPassword());
return new JdbcTemplate(ds);
}
private static void verifyRequestHasBeenAdvisedWithMessagesFromVectorStore(ChatModel chatModel) {
ArgumentCaptor<Prompt> promptCaptor = ArgumentCaptor.forClass(Prompt.class);
verify(chatModel).call(promptCaptor.capture());
assertThat(promptCaptor.getValue().getInstructions().get(0)).isInstanceOf(SystemMessage.class);
assertThat(promptCaptor.getValue().getInstructions().get(0).getContent()).isEqualTo("""
Use the long term conversation memory from the LONG_TERM_MEMORY section to provide accurate answers.
---------------------
LONG_TERM_MEMORY:
Tell me a good joke
Tell me a bad joke
---------------------
""");
}
/**
* Test that chats with {@link VectorStoreChatMemoryAdvisor} get advised with similar
* messages from the (gp)vector store.
@@ -88,42 +143,6 @@ class PgVectorStoreWithChatMemoryAdvisorIT {
verifyRequestHasBeenAdvisedWithMessagesFromVectorStore(chatModel);
}
private static @NotNull ChatModel chatModelAlwaysReturnsTheSameReply() {
ChatModel chatModel = mock(ChatModel.class);
ArgumentCaptor<Prompt> argumentCaptor = ArgumentCaptor.forClass(Prompt.class);
ChatResponse chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
Why don't scientists trust atoms?
Because they make up everything!
"""))));
when(chatModel.call(argumentCaptor.capture())).thenReturn(chatResponse);
return chatModel;
}
private static void initStore(PgVectorStore store) throws Exception {
store.afterPropertiesSet();
// fill the store
store.add(List.of(new Document("Tell me a good joke", Map.of("conversationId", "default")),
new Document("Tell me a bad joke", Map.of("conversationId", "default", "messageType", "USER"))));
}
private static PgVectorStore createPgVectorStoreUsingTestcontainer(EmbeddingModel embeddingModel) throws Exception {
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
PgVectorStore vectorStore = new PgVectorStore.Builder(jdbcTemplate, embeddingModel).withDimensions(3) // match
// embeddings
.withInitializeSchema(true)
.build();
initStore(vectorStore);
return vectorStore;
}
private static @NotNull JdbcTemplate createJdbcTemplateWithConnectionToTestcontainer() {
PGSimpleDataSource ds = new PGSimpleDataSource();
ds.setUrl("jdbc:postgresql://localhost:" + postgresContainer.getMappedPort(5432) + "/postgres");
ds.setUser(postgresContainer.getUsername());
ds.setPassword(postgresContainer.getPassword());
return new JdbcTemplate(ds);
}
@SuppressWarnings("unchecked")
private @NotNull EmbeddingModel embeddingNModelShouldAlwaysReturnFakedEmbed() {
EmbeddingModel embeddingModel = mock(EmbeddingModel.class);
@@ -131,29 +150,11 @@ class PgVectorStoreWithChatMemoryAdvisorIT {
Mockito.doAnswer(invocationOnMock -> {
Object[] arguments = invocationOnMock.getArguments();
List<Document> documents = (List<Document>) arguments[0];
documents.forEach(d -> d.setEmbedding(embed));
return List.of(embed, embed);
documents.forEach(d -> d.setEmbedding(this.embed));
return List.of(this.embed, this.embed);
}).when(embeddingModel).embed(ArgumentMatchers.any(), any(), any());
when(embeddingModel.embed(any(String.class))).thenReturn(embed);
when(embeddingModel.embed(any(String.class))).thenReturn(this.embed);
return embeddingModel;
}
private static void verifyRequestHasBeenAdvisedWithMessagesFromVectorStore(ChatModel chatModel) {
ArgumentCaptor<Prompt> promptCaptor = ArgumentCaptor.forClass(Prompt.class);
verify(chatModel).call(promptCaptor.capture());
assertThat(promptCaptor.getValue().getInstructions().get(0)).isInstanceOf(SystemMessage.class);
assertThat(promptCaptor.getValue().getInstructions().get(0).getContent()).isEqualTo("""
Use the long term conversation memory from the LONG_TERM_MEMORY section to provide accurate answers.
---------------------
LONG_TERM_MEMORY:
Tell me a good joke
Tell me a bad joke
---------------------
""");
}
}
}