From 1d479e3aabbd88ecf4ce90a80b4a8215368b4fbb Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sat, 15 Jun 2024 17:49:10 +0200 Subject: [PATCH] Minor PgVector ITs improvements --- .../PgVectorStoreAutoConfiguration2IT.java | 114 ++++++++++++++++++ .../PgVectorStoreAutoConfigurationIT.java | 2 +- .../spring-ai-pgvector-store/pom.xml | 5 + .../ai/vectorstore/PgVectorStoreIT.java | 22 ++-- 4 files changed, 131 insertions(+), 12 deletions(-) create mode 100644 spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java new file mode 100644 index 000000000..c31a59368 --- /dev/null +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfiguration2IT.java @@ -0,0 +1,114 @@ +/* + * 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 + * + * 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.autoconfigure.vectorstore.pgvector; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.Map; + +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; +import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.io.DefaultResourceLoader; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +@Testcontainers +public class PgVectorStoreAutoConfiguration2IT { + + @Container + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16") + .withUsername("postgres") + .withPassword("postgres"); + + List documents = List.of( + new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("spring", "great")), + new Document(getText("classpath:/test/data/time.shelter.txt")), + new Document(getText("classpath:/test/data/great.depression.txt"), Map.of("depression", "bad"))); + + 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() + .withConfiguration(AutoConfigurations.of(PgVectorStoreAutoConfiguration.class, + JdbcTemplateAutoConfiguration.class, DataSourceAutoConfiguration.class)) + .withUserConfiguration(Config.class) + .withPropertyValues("spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE", + // JdbcTemplate configuration + String.format("spring.datasource.url=jdbc:postgresql://%s:%d/%s", postgresContainer.getHost(), + postgresContainer.getMappedPort(5432), "postgres"), + "spring.datasource.username=postgres", "spring.datasource.password=postgres"); + + @Test + public void addAndSearch() { + contextRunner.run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + vectorStore.add(documents); + + List 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("depression", "distance"); + + // Remove all documents from the store + vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList()); + results = vectorStore.similaritySearch(SearchRequest.query("Great Depression").withTopK(1)); + assertThat(results).hasSize(0); + }); + } + + @Configuration(proxyBeanMethods = false) + static class Config { + + @Bean + public EmbeddingModel embeddingModel() { + return new TransformersEmbeddingModel(); + } + + } + +} diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java index d3cbf9575..346145760 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/vectorstore/pgvector/PgVectorStoreAutoConfigurationIT.java @@ -47,7 +47,7 @@ import static org.assertj.core.api.Assertions.assertThat; public class PgVectorStoreAutoConfigurationIT { @Container - static GenericContainer postgresContainer = new GenericContainer<>("ankane/pgvector:v0.5.1") + static GenericContainer postgresContainer = new GenericContainer<>("pgvector/pgvector:0.7.2-pg16") .withEnv("POSTGRES_USER", "postgres") .withEnv("POSTGRES_PASSWORD", "postgres") .withExposedPorts(5432); diff --git a/vector-stores/spring-ai-pgvector-store/pom.xml b/vector-stores/spring-ai-pgvector-store/pom.xml index 056e6e172..0da1ef575 100644 --- a/vector-stores/spring-ai-pgvector-store/pom.xml +++ b/vector-stores/spring-ai-pgvector-store/pom.xml @@ -76,6 +76,11 @@ testcontainers test + + org.testcontainers + postgresql + test + org.testcontainers diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java index 8ac1219ff..19586503d 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java @@ -15,6 +15,8 @@ */ 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.Collections; @@ -25,19 +27,14 @@ import java.util.UUID; 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.ValueSource; -import org.testcontainers.containers.GenericContainer; -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.api.OpenAiApi; 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.filter.FilterExpressionTextParser.FilterExpressionParseException; import org.springframework.beans.factory.annotation.Value; @@ -53,8 +50,11 @@ 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 static org.assertj.core.api.Assertions.assertThat; +import com.zaxxer.hikari.HikariDataSource; /** * @author Christian Tzolov @@ -64,10 +64,10 @@ import static org.assertj.core.api.Assertions.assertThat; public class PgVectorStoreIT { @Container - static GenericContainer postgresContainer = new GenericContainer<>("ankane/pgvector:v0.5.1") - .withEnv("POSTGRES_USER", "postgres") - .withEnv("POSTGRES_PASSWORD", "postgres") - .withExposedPorts(5432); + @SuppressWarnings("resource") + static PostgreSQLContainer postgresContainer = new PostgreSQLContainer<>("pgvector/pgvector:pg16") + .withUsername("postgres") + .withPassword("postgres"); List documents = List.of( new Document(getText("classpath:/test/data/spring.ai.txt"), Map.of("meta1", "meta1")),