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

Request-time filter expressions for RAG

Browse Source 
When using the RetrievalAugmentationAdvisor with the VectorStoreDocumentRetriever, it’s now possible to provide a filter expression at request-time as an advisor context variable with key VectorStoreDocumentRetriever.FILTER_EXPRESSION.

Fixes gh-1776

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2025-03-09 11:29:34 +01:00
committed by Ilayaperumal Gopinathan
parent 82b46d2182
commit 5a4e9f5108
6 changed files with 132 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java
View File

@@ -108,6 +108,29 @@ class RetrievalAugmentationAdvisorIT {
evaluateRelevancy(question, chatResponse);
}
@Test
void ragWithRequestFilter() {
String question = "Where does the adventure of Anacletus and Birba take place?";
RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
.documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
.build();
ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
.build()
.prompt(question)
.advisors(ragAdvisor)
.advisors(a -> a.param(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Italy'"))
.call()
.chatResponse();
assertThat(chatResponse).isNotNull();
// No documents retrieved since the filter expression matches none of the
// documents in the vector store.
assertThat((String) chatResponse.getResult().getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT))
.isNull();
}
@Test
void ragWithCompression() {
MessageChatMemoryAdvisor memoryAdvisor = MessageChatMemoryAdvisor.builder(new InMemoryChatMemory()).build();

4
spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/preretrieval/query/transformation/RewriteQueryTransformerIT.java
View File

@@ -43,7 +43,7 @@ class RewriteQueryTransformerIT {
@Test
void whenTransformerWithDefaults() {
Query query = new Query("I'm studying machine learning. What is an LLM?");
Query query = new Query("What are the main tourist attractions in L.A.?");
QueryTransformer queryTransformer = RewriteQueryTransformer.builder()
.chatClientBuilder(ChatClient.builder(this.openAiChatModel))
.build();
@@ -52,7 +52,7 @@ class RewriteQueryTransformerIT {
assertThat(transformedQuery).isNotNull();
System.out.println(transformedQuery);
assertThat(transformedQuery.text()).containsIgnoringCase("model");
assertThat(transformedQuery.text()).containsIgnoringCase("Angeles");
}
}

44
spring-ai-integration-tests/src/test/java/org/springframework/ai/integration/tests/rag/retrieval/search/VectorStoreDocumentRetrieverIT.java
View File

@@ -46,20 +46,21 @@ import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.EQ
@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(
// @formatter:off
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(
"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(
"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")));
// @formatter:on
@Autowired
PgVectorStore pgVectorStore;
@@ -75,7 +76,7 @@ class VectorStoreDocumentRetrieverIT {
}
@Test
void withFilter() {
void withBuildFilter() {
DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
.vectorStore(this.pgVectorStore)
.similarityThreshold(0.50)
@@ -95,7 +96,7 @@ class VectorStoreDocumentRetrieverIT {
}
@Test
void withNoFilter() {
void withNoBuildFilter() {
DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
.vectorStore(this.pgVectorStore)
.similarityThreshold(0.50)
@@ -110,4 +111,27 @@ class VectorStoreDocumentRetrieverIT {
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("3").getId()));
}
@Test
void withRequestFilter() {
DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
.vectorStore(this.pgVectorStore)
.similarityThreshold(0.50)
.topK(3)
.build();
Query query = Query.builder()
.text("Who is Anacletus?")
.context(Map.of(VectorStoreDocumentRetriever.FILTER_EXPRESSION, "location == 'Whispering Woods'"))
.build();
List<Document> retrievedDocuments = documentRetriever.retrieve(query);
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()));
// No request filter expression applied, so full access to all documents.
retrievedDocuments = documentRetriever.retrieve(new Query("Who is Birba?"));
assertThat(retrievedDocuments).anyMatch(document -> document.getId().equals(documents.get("4").getId()));
}
}