Modular RAG - Query Analysis
Query Analysis * Introduce Query Analysis Module * Define QueryTransformer API and TranslationQueryTransformer implementation * Define QueryExpander API and MultiQueryExpander implementation * Support QueryTransformer in RetrievalAugmentationAdvisor (support for QueryExpander will be in the next PR together with the needed DocumentFuser API). Improvements * Refine Retrieval and Augmentation Modules for increased robustness * Expand test coverage for both modules * Define clone() method for ChatClient.Builder Tests * Introduce “spring-ai-integration-tests” for full-fledged integration tests * Add integration tests for RAG modules * Add integration tests for RAG advisor Query Analysis * Introduce Query Analysis Module * Define QueryTransformer API and TranslationQueryTransformer implementation * Define QueryExpander API and MultiQueryExpander implementation * Support QueryTransformer in RetrievalAugmentationAdvisor (support for QueryExpander will be in the next PR together with the needed DocumentFuser API). Improvements * Refine Retrieval and Augmentation Modules for increased robustness * Expand test coverage for both modules * Define clone() method for ChatClient.Builder Tests * Introduce “spring-ai-integration-tests” for full-fledged integration tests * Add integration tests for RAG modules * Add integration tests for RAG advisor Relates to #gh-1603 Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
committed by
Mark Pollack
Mark Pollack
parent
b4e0a4598e
commit
263fe2fba7
@@ -0,0 +1,30 @@
|
||||
/*
|
||||
* 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.integration.tests;
|
||||
|
||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||
import org.springframework.context.annotation.Import;
|
||||
|
||||
/**
|
||||
* Test application for integration tests.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootApplication
|
||||
@Import(TestcontainersConfiguration.class)
|
||||
public class TestApplication {
|
||||
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
/*
|
||||
* 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.integration.tests;
|
||||
|
||||
import org.springframework.boot.test.context.TestConfiguration;
|
||||
import org.springframework.boot.testcontainers.service.connection.ServiceConnection;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.testcontainers.containers.PostgreSQLContainer;
|
||||
|
||||
/**
|
||||
* Test configuration for Testcontainers-based Dev Services.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@TestConfiguration(proxyBeanMethods = false)
|
||||
class TestcontainersConfiguration {
|
||||
|
||||
@Bean
|
||||
@ServiceConnection
|
||||
PostgreSQLContainer<?> pgvectorContainer() {
|
||||
return new PostgreSQLContainer<>("pgvector/pgvector:pg17");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
/*
|
||||
* 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.integration.tests.client.advisor;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentReader;
|
||||
import org.springframework.ai.evaluation.EvaluationRequest;
|
||||
import org.springframework.ai.evaluation.EvaluationResponse;
|
||||
import org.springframework.ai.evaluation.RelevancyEvaluator;
|
||||
import org.springframework.ai.integration.tests.TestApplication;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.rag.analysis.query.transformation.TranslationQueryTransformer;
|
||||
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
|
||||
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
|
||||
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
|
||||
import org.springframework.ai.vectorstore.PgVectorStore;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.core.io.Resource;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link RetrievalAugmentationAdvisor}.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootTest(classes = TestApplication.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
|
||||
class RetrievalAugmentationAdvisorIT {
|
||||
|
||||
private List<Document> knowledgeBaseDocuments;
|
||||
|
||||
@Autowired
|
||||
OpenAiChatModel openAiChatModel;
|
||||
|
||||
@Autowired
|
||||
PgVectorStore pgVectorStore;
|
||||
|
||||
@Value("${classpath:documents/knowledge-base.md}")
|
||||
Resource knowledgeBaseResource;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
DocumentReader markdownReader = new MarkdownDocumentReader(knowledgeBaseResource,
|
||||
MarkdownDocumentReaderConfig.defaultConfig());
|
||||
knowledgeBaseDocuments = markdownReader.read();
|
||||
pgVectorStore.add(knowledgeBaseDocuments);
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
pgVectorStore.delete(knowledgeBaseDocuments.stream().map(Document::getId).toList());
|
||||
}
|
||||
|
||||
@Test
|
||||
void ragBasic() {
|
||||
String question = "Where does the adventure of Anacletus and Birba take place?";
|
||||
|
||||
RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
|
||||
.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(pgVectorStore).build())
|
||||
.build();
|
||||
|
||||
ChatResponse chatResponse = ChatClient.builder(openAiChatModel)
|
||||
.build()
|
||||
.prompt(question)
|
||||
.advisors(ragAdvisor)
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
assertThat(chatResponse).isNotNull();
|
||||
|
||||
String response = chatResponse.getResult().getOutput().getContent();
|
||||
System.out.println(response);
|
||||
assertThat(response).containsIgnoringCase("Highlands");
|
||||
|
||||
evaluateRelevancy(question, chatResponse);
|
||||
}
|
||||
|
||||
@Test
|
||||
void ragWithTranslation() {
|
||||
String question = "Hvor finder Anacletus og Birbas eventyr sted?";
|
||||
|
||||
RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
|
||||
.queryTransformers(TranslationQueryTransformer.builder()
|
||||
.chatClientBuilder(ChatClient.builder(openAiChatModel))
|
||||
.targetLanguage("english")
|
||||
.build())
|
||||
.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(pgVectorStore).build())
|
||||
.build();
|
||||
|
||||
ChatResponse chatResponse = ChatClient.builder(openAiChatModel)
|
||||
.build()
|
||||
.prompt(question)
|
||||
.advisors(ragAdvisor)
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
assertThat(chatResponse).isNotNull();
|
||||
|
||||
String response = chatResponse.getResult().getOutput().getContent();
|
||||
System.out.println(response);
|
||||
assertThat(response).containsIgnoringCase("Highlands");
|
||||
|
||||
evaluateRelevancy(question, chatResponse);
|
||||
}
|
||||
|
||||
private void evaluateRelevancy(String question, ChatResponse chatResponse) {
|
||||
EvaluationRequest evaluationRequest = new EvaluationRequest(question,
|
||||
chatResponse.getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT),
|
||||
chatResponse.getResult().getOutput().getContent());
|
||||
RelevancyEvaluator evaluator = new RelevancyEvaluator(ChatClient.builder(openAiChatModel));
|
||||
EvaluationResponse evaluationResponse = evaluator.evaluate(evaluationRequest);
|
||||
assertThat(evaluationResponse.isPass()).isTrue();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
/*
|
||||
* 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.integration.tests.rag.analysis.query.expansion;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.integration.tests.TestApplication;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.analysis.query.expansion.MultiQueryExpander;
|
||||
import org.springframework.ai.rag.analysis.query.expansion.QueryExpander;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link MultiQueryExpander}.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootTest(classes = TestApplication.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
|
||||
class MultiQueryExpanderIT {
|
||||
|
||||
@Autowired
|
||||
OpenAiChatModel openAiChatModel;
|
||||
|
||||
@Test
|
||||
void whenExpanderWithDefaults() {
|
||||
Query query = new Query("What is the weather in Rome?");
|
||||
QueryExpander queryExpander = MultiQueryExpander.builder()
|
||||
.chatClientBuilder(ChatClient.builder(openAiChatModel))
|
||||
.build();
|
||||
|
||||
List<Query> queries = queryExpander.apply(query);
|
||||
|
||||
assertThat(queries).isNotNull();
|
||||
queries.forEach(System.out::println);
|
||||
assertThat(queries).hasSize(3);
|
||||
}
|
||||
|
||||
@Test
|
||||
void whenExpanderWithCustomQueryNumber() {
|
||||
Query query = new Query("What is the weather in Rome?");
|
||||
QueryExpander queryExpander = MultiQueryExpander.builder()
|
||||
.chatClientBuilder(ChatClient.builder(openAiChatModel))
|
||||
.numberOfQueries(4)
|
||||
.build();
|
||||
|
||||
List<Query> queries = queryExpander.apply(query);
|
||||
|
||||
assertThat(queries).isNotNull();
|
||||
queries.forEach(System.out::println);
|
||||
assertThat(queries).hasSize(4);
|
||||
}
|
||||
|
||||
@Test
|
||||
void whenExpanderWithOriginalQueryIncluded() {
|
||||
Query query = new Query("What is the weather in Rome?");
|
||||
QueryExpander queryExpander = MultiQueryExpander.builder()
|
||||
.chatClientBuilder(ChatClient.builder(openAiChatModel))
|
||||
.numberOfQueries(3)
|
||||
.includeOriginal(true)
|
||||
.build();
|
||||
|
||||
List<Query> queries = queryExpander.apply(query);
|
||||
|
||||
assertThat(queries).isNotNull();
|
||||
queries.forEach(System.out::println);
|
||||
assertThat(queries).hasSize(4);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
/*
|
||||
* 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.integration.tests.rag.analysis.query.transformation;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.integration.tests.TestApplication;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.analysis.query.transformation.QueryTransformer;
|
||||
import org.springframework.ai.rag.analysis.query.transformation.TranslationQueryTransformer;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link TranslationQueryTransformer}.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootTest(classes = TestApplication.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
|
||||
class TranslationQueryTransformerIT {
|
||||
|
||||
@Autowired
|
||||
OpenAiChatModel openAiChatModel;
|
||||
|
||||
@Test
|
||||
void whenTransformerWithDefaults() {
|
||||
Query query = new Query("Hvad er Danmarks hovedstad?");
|
||||
QueryTransformer queryTransformer = TranslationQueryTransformer.builder()
|
||||
.chatClientBuilder(ChatClient.builder(openAiChatModel))
|
||||
.targetLanguage("english")
|
||||
.build();
|
||||
|
||||
Query transformedQuery = queryTransformer.apply(query);
|
||||
|
||||
assertThat(transformedQuery).isNotNull();
|
||||
System.out.println(transformedQuery);
|
||||
assertThat(transformedQuery.text()).containsIgnoringCase("Denmark").containsIgnoringCase("capital");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
/*
|
||||
* 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.integration.tests.rag.augmentation;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.integration.tests.TestApplication;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.augmentation.ContextualQueryAugmentor;
|
||||
import org.springframework.ai.rag.augmentation.QueryAugmentor;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link ContextualQueryAugmentor}.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootTest(classes = TestApplication.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
|
||||
class ContextualQueryAugmentorIT {
|
||||
|
||||
@Autowired
|
||||
OpenAiChatModel openAiChatModel;
|
||||
|
||||
@Test
|
||||
void whenContextIsProvided() {
|
||||
QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().build();
|
||||
Query query = new Query("What is Iorek's dream?");
|
||||
List<Document> documents = List
|
||||
.of(new Document("Iorek was a little polar bear who lived in the Arctic circle."), new Document(
|
||||
"Iorek loved to explore the snowy landscape and dreamt of one day going on an adventure around the North Pole."));
|
||||
|
||||
Query augmentedQuery = queryAugmentor.augment(query, documents);
|
||||
String response = openAiChatModel.call(augmentedQuery.text());
|
||||
|
||||
assertThat(response).isNotEmpty();
|
||||
System.out.println(response);
|
||||
assertThat(response).containsIgnoringCase("North Pole");
|
||||
assertThat(response).doesNotContainIgnoringCase("context");
|
||||
assertThat(response).doesNotContainIgnoringCase("information");
|
||||
}
|
||||
|
||||
@Test
|
||||
void whenAllowEmptyContext() {
|
||||
QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().build();
|
||||
Query query = new Query("What is Iorek's dream?");
|
||||
List<Document> documents = List.of();
|
||||
Query augmentedQuery = queryAugmentor.augment(query, documents);
|
||||
String response = openAiChatModel.call(augmentedQuery.text());
|
||||
|
||||
assertThat(response).isNotEmpty();
|
||||
System.out.println(response);
|
||||
assertThat(response).containsIgnoringCase("Iorek");
|
||||
}
|
||||
|
||||
@Test
|
||||
void whenNotAllowEmptyContext() {
|
||||
QueryAugmentor queryAugmentor = ContextualQueryAugmentor.builder().allowEmptyContext(false).build();
|
||||
Query query = new Query("What is Iorek's dream?");
|
||||
List<Document> documents = List.of();
|
||||
Query augmentedQuery = queryAugmentor.augment(query, documents);
|
||||
String response = openAiChatModel.call(augmentedQuery.text());
|
||||
|
||||
assertThat(response).isNotEmpty();
|
||||
System.out.println(response);
|
||||
assertThat(response).doesNotContainIgnoringCase("Iorek");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
/*
|
||||
* 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.integration.tests.rag.retrieval.search;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.integration.tests.TestApplication;
|
||||
import org.springframework.ai.rag.Query;
|
||||
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
|
||||
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
|
||||
import org.springframework.ai.vectorstore.PgVectorStore;
|
||||
import org.springframework.ai.vectorstore.filter.Filter;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link VectorStoreDocumentRetriever}.
|
||||
*
|
||||
* @author Thomas Vitale
|
||||
*/
|
||||
@SpringBootTest(classes = TestApplication.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".*")
|
||||
class VectorStoreDocumentRetrieverIT {
|
||||
|
||||
private static final Map<String, Document> documents = Map.of("1", new Document(
|
||||
"Anacletus was a majestic snowy owl with unusually bright golden eyes and distinctive black speckles across his wings.",
|
||||
Map.of("location", "Whispering Woods")), "2",
|
||||
new Document(
|
||||
"Anacletus made his home in an ancient hollow oak tree deep within the Whispering Woods, where local villagers often heard his haunting calls at midnight.",
|
||||
Map.of("location", "Whispering Woods")),
|
||||
"3",
|
||||
new Document(
|
||||
"Despite being a nocturnal hunter like other owls, Anacletus had developed a peculiar habit of collecting shiny objects, especially lost coins and jewelry that glinted in the moonlight.",
|
||||
Map.of()),
|
||||
"4",
|
||||
new Document(
|
||||
"Birba was a plump Siamese cat with mismatched eyes - one blue and one green - who spent her days lounging on velvet cushions and judging everyone with a perpetual look of disdain.",
|
||||
Map.of("location", "Alfea")));
|
||||
|
||||
@Autowired
|
||||
PgVectorStore pgVectorStore;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
pgVectorStore.add(List.copyOf(documents.values()));
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
pgVectorStore.delete(documents.values().stream().map(Document::getId).toList());
|
||||
}
|
||||
|
||||
@Test
|
||||
void withFilter() {
|
||||
DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
|
||||
.vectorStore(pgVectorStore)
|
||||
.similarityThreshold(0.50)
|
||||
.topK(3)
|
||||
.filterExpression(
|
||||
new Filter.Expression(EQ, new Filter.Key("location"), new Filter.Value("Whispering Woods")))
|
||||
.build();
|
||||
|
||||
List<Document> retrievedDocuments = documentRetriever.retrieve(new Query("Who is Anacletus?"));
|
||||
|
||||
assertThat(retrievedDocuments).hasSize(2);
|
||||
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("1").getId()));
|
||||
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("2").getId()));
|
||||
|
||||
retrievedDocuments = documentRetriever.retrieve(new Query("Who is Birba?"));
|
||||
assertThat(retrievedDocuments).noneMatch(document -> document.getId().equals(documents.get("4").getId()));
|
||||
}
|
||||
|
||||
@Test
|
||||
void withNoFilter() {
|
||||
DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
|
||||
.vectorStore(pgVectorStore)
|
||||
.similarityThreshold(0.50)
|
||||
.topK(3)
|
||||
.build();
|
||||
|
||||
List<Document> retrievedDocuments = documentRetriever.retrieve(new Query("Who is Anacletus?"));
|
||||
|
||||
assertThat(retrievedDocuments).hasSize(3);
|
||||
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("1").getId()));
|
||||
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("2").getId()));
|
||||
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("3").getId()));
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
spring:
|
||||
main:
|
||||
web-application-type: none
|
||||
ai:
|
||||
openai:
|
||||
api-key: ${OPENAI_API_KEY}
|
||||
chat:
|
||||
options:
|
||||
model: gpt-4o-mini
|
||||
embedding:
|
||||
options:
|
||||
model: text-embedding-ada-002
|
||||
retry:
|
||||
max-attempts: 3
|
||||
vectorstore:
|
||||
pgvector:
|
||||
initialize-schema: true
|
||||
@@ -0,0 +1,41 @@
|
||||
# Anacletus and Birba's Quest for the Loch of the Stars
|
||||
|
||||
## Chapter 1: The Map and the Adventure
|
||||
|
||||
Once upon a time, in a cozy little cottage nestled at the edge of the Scottish Highlands, lived an owl named Anacletus and a curious cat named Birba. Anacletus was wise and careful, always reading maps and planning things thoroughly, while Birba was lively and adventurous, always ready to chase after the next interesting thing. Despite their differences, they were the best of friends and loved going on little adventures together.
|
||||
|
||||
## Chapter 2: The Journey Begins
|
||||
|
||||
One sunny morning, Anacletus showed Birba an old, crinkled map he’d found in the attic. “Look, Birba,” he said, pointing with his feathery wing. “This map leads to the legendary Loch of the Stars. They say it shines brighter than any other lake at night.” Birba’s eyes sparkled with excitement. “Oh, we have to go there!” she meowed. So, they packed a small bag with snacks, a compass, and a flashlight, and off they went, eager to find the legendary loch.
|
||||
|
||||
## Chapter 3: The Highland Adventure
|
||||
|
||||
Their journey began with a climb up the rolling hills covered in purple heather. Anacletus flapped his wings, soaring ahead to scout for any obstacles, while Birba trotted along below, her nose sniffing the air for interesting scents. Soon, they came across a bubbling brook. Anacletus carefully flew over it, but Birba hesitated. “Just a little jump!” Anacletus called out. With a deep breath, Birba leaped and landed safely on the other side. She purred proudly, and they continued on their way.
|
||||
|
||||
## Chapter 4: The Highland Cows and the Hidden Path
|
||||
|
||||
As they ventured deeper into the Highlands, they stumbled upon a herd of curious Highland cows with long, shaggy hair. The cows mooed softly, and one of them named Fergus approached. “Where are you two headed?” Fergus asked. “We’re searching for the Loch of the Stars!” Anacletus replied. Fergus nodded knowingly and pointed his nose north. “Follow the path by the big stones, and it will lead you closer to the loch,” he said. Thanking Fergus, they set off again, Birba occasionally stopping to bat at the fluttering butterflies along the way.
|
||||
|
||||
## Chapter 5: The Mysterious Forest and the Deer Family
|
||||
|
||||
The day wore on, and they soon found themselves in a mysterious forest. Tall, ancient pine trees surrounded them, casting long shadows. “Stay close, Birba,” Anacletus whispered, his wise eyes scanning for any sign of danger. But Birba had already darted after a flicker of light, thinking it was a firefly. Anacletus sighed and followed her until they came to a hidden glade where a family of deer grazed quietly. The smallest fawn looked up and gave them a curious nod before they moved along.
|
||||
|
||||
## Chapter 6: The Loch of the Stars
|
||||
|
||||
After a while, the sun began to set, painting the sky in shades of pink and gold. Anacletus decided it was a good time to rest. They found a cozy hollow at the base of a tree, where they shared the snacks they’d packed. Birba munched on her fish treats while Anacletus nibbled on a biscuit. “Do you think we’ll find the Loch of the Stars?” Birba asked, her eyes twinkling. “I think so,” Anacletus replied with a wise smile. “We’re getting closer.”
|
||||
|
||||
## Chapter 7: The Shimmering Loch
|
||||
|
||||
As night fell, they finally reached the top of a hill where they could see a shimmering light in the distance. “Look, Birba!” Anacletus hooted excitedly. There, nestled among the hills, was the Loch of the Stars, gleaming like a sky full of stars. The two friends hurried down to the water’s edge, marveling at how the loch sparkled under the moonlight, casting a gentle glow all around.
|
||||
|
||||
## Chapter 8: The Magic of the Loch
|
||||
|
||||
Birba dipped a curious paw into the water, causing ripples that sent stars dancing across the surface. “It’s beautiful!” she gasped. Anacletus nodded, his heart filled with awe. They spent the night by the loch, watching the shimmering stars reflected in the water, feeling as though they were surrounded by magic.
|
||||
|
||||
## Chapter 9: The Journey Home
|
||||
|
||||
When dawn broke, the shimmering loch returned to its quiet, glassy calm. With a satisfied yawn, Birba stretched and said, “That was the best adventure yet.” Anacletus agreed, feeling a warmth in his feathers as they turned back toward home, carrying memories of the Loch of the Stars in their hearts.
|
||||
|
||||
## Chapter 10: The End of the Adventure
|
||||
|
||||
And as they made their way back to their cozy cottage, they already started dreaming of their next big adventure—because Anacletus and Birba knew that the Scottish Highlands held endless wonders for those who dared to explore.
|
||||
Reference in New Issue
Block a user