Modular RAG - Query Analysis #1703

ThomasVitale · 2024-11-09T08:51:29Z

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

markpollack · 2024-11-13T00:21:08Z

spring-ai-core/src/main/java/org/springframework/ai/chat/client/DefaultChatClientBuilder.java

 		return new DefaultChatClient(this.defaultRequest);
 	}

+	public Builder clone() {


All the other APIs used by ChatClient have similar methods, but they are called mutate() instead of clone(). I wonder if it'd make sense to change them to clone()? It would be aligned with the naming used in other Spring APIs (like RestClient) and also aligned with Java conventions in general.

We borrowed the mutate from RestClient : https://github.com/spring-projects/spring-framework/blob/b523f3caff8b57feca7b445a1a151cb9aa1f7fb3/spring-web/src/main/java/org/springframework/web/client/RestClient.java#L136
So mutate is the Spring API style

As discussed, we figured the mystery :) The Builder has the "clone()" method and the Client APIs the "mutate()" method.

markpollack · 2024-11-13T00:34:16Z

...e/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

 		return AdvisedRequest.from(request).withUserText(augmentedQuery.text()).withAdviseContext(context).build();
 	}

 	private AdvisedResponse after(AdvisedResponse advisedResponse) {


could make protected, along with before if someone wants to customize the algorithm.

As discussed, the customisation approach will be based on implementing a custom Advisor and using the available RAG components to build a custom RAG flow, whereas all these out-of-the-box implementations will be "final" following the Spring Security approach.

markpollack · 2024-11-13T00:46:15Z

Just some general thoughts, not necessary to resolve for this PR.

The Query record. This is essentially the user text. What is this was instead the UserMessage so that it can take into account multi-modality (media) A QueryTransformer could for example retrieve images related to the text.
The QueryAugmentor - what if we had as the signature QueryAugmentor extends BiFunction<Query, RetrievalContext, Query> where RetrievalContext is

public class RetrievalContext {
      private final List<Document> documents;
      private final Map<String, Object> metadata;
   ...

this way additional information that can be used by an implementing class can be taken into account. The document is the core part, but there could be additional information to influence how the query gets augmented. Perhaps that would be passed in the constrcutor to the QueryAugmentor impl. Just looking to move away from types that don't have extensibility in the interface.

Similar for DocumentRetriever, instead of List it returns a class RetrievedDocuments that contains the list plus additional data about those documents?

markpollack · 2024-11-13T00:53:21Z

...e/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

+		// 3. Augment user query with the document contextual data.
+		Query augmentedQuery = this.queryAugmentor.augment(transformedQuery, documents);

 		return AdvisedRequest.from(request).withUserText(augmentedQuery.text()).withAdviseContext(context).build();


wondering about the multimodal/media aspects of this as it focused on text and any media for the request is lost.

The AdvisedRequest has a separate Media property for the media elements, so they are currently propagated correctly.

markpollack · 2024-11-13T00:56:39Z

...e/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

-		List<Document> documents = this.documentRetriever.retrieve(query);
+		// 2. Retrieve similar documents for the original query.
+		List<Document> documents = this.documentRetriever.retrieve(transformedQuery);
 		context.put(DOCUMENT_CONTEXT, documents);


maybe a RAGContext that in turn contains the documents so that for the future when additional things are added to the context for RAG, it is all in a central spot?

We agreed to talk more about this when we reconsider the design of ChatResponse and how to propagate back evidence from different intermediate steps in a ChatClient, such as retrieved documents and called functions.

markpollack · 2024-11-13T01:02:37Z

...e/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

+		this.queryTransformers = queryTransformers;
 		this.documentRetriever = documentRetriever;
 		this.queryAugmentor = queryAugmentor != null ? queryAugmentor : ContextualQueryAugmentor.builder().build();
 		this.protectFromBlocking = protectFromBlocking != null ? protectFromBlocking : false;


I think protectFromBlocking being true by default is a more conservative option, this way it works out of the box in a reactive environment and for a non-reactive environment there isn't much of an impact.

Good catch! It's now true by default.

markpollack · 2024-11-13T01:05:54Z

...ai-core/src/main/java/org/springframework/ai/rag/analysis/query/expansion/QueryExpander.java

+ * @author Thomas Vitale
+ * @since 1.0.0
+ */
+public interface QueryExpander extends Function<Query, List<Query>> {


i'm wondering if this is something that should be used inside RetrievalAugmentationAdvisor or it is used in code at higher level. Maybe this is part of the Query Optimization part that is to come

I haven't adopted that in RetrievalAugmentationAdvisor yet, but I'll do so in a next PR where I introduce the Router and Fuser components that are needed to handle multiple queries/retrievals.

markpollack · 2024-11-13T01:18:00Z

.../org/springframework/ai/integration/tests/client/advisor/RetrievalAugmentationAdvisorIT.java

+	private List<Document> knowledgeBaseDocuments;
+
+	@Autowired
+	OpenAiChatModel openAiChatModel;


I'm getting a test failure

Caused by: java.lang.IllegalArgumentException: OpenAI API key must be set. Use the connection property: spring.ai.openai.api-key or spring.ai.openai.chat.api-key property. at org.springframework.util.Assert.hasText(Assert.java:240) ~[spring-core-6.1.13.jar:6.1.13] at org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration.resolveConnectionProperties(OpenAiAutoConfiguration.java:100) ~[classes/:na]

I think some additional context config is needed.

I'll look into that. I wonder if it's due to the new APIKey concept introduced after I opened this PR

Oh yeah, it's actually the missing property you mentioned in #1703 (comment) I'll fix that

markpollack · 2024-11-13T01:20:39Z

I get test failures for tests org.springframework.ai.integration.tests due to the api-key not being set. Prob works in your environment but in my I don't have any spring.ai.openai.*properties set.

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 #spring-projectsgh-1603 Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>

markpollack · 2024-11-13T21:50:09Z

merged in 263fe2f

ThomasVitale mentioned this pull request Nov 9, 2024

Modular RAG - Part 1 #1603

Closed

ThomasVitale force-pushed the modular-rag-query-analysis branch 2 times, most recently from 90f6791 to 69494ec Compare November 10, 2024 13:11

markpollack self-assigned this Nov 11, 2024

tzolov self-assigned this Nov 12, 2024

markpollack reviewed Nov 13, 2024

View reviewed changes

ThomasVitale added 2 commits November 13, 2024 20:23

ThomasVitale force-pushed the modular-rag-query-analysis branch from 69494ec to f902229 Compare November 13, 2024 19:24

markpollack closed this Nov 13, 2024

markpollack added this to the 1.0.0-M4 milestone Nov 14, 2024

Modular RAG - Query Analysis #1703

Modular RAG - Query Analysis #1703

Uh oh!

Conversation

ThomasVitale commented Nov 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markpollack commented Nov 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markpollack Nov 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markpollack Nov 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

markpollack commented Nov 13, 2024

Uh oh!

markpollack commented Nov 13, 2024

Uh oh!

Uh oh!

ThomasVitale commented Nov 9, 2024 •

edited

Loading

markpollack commented Nov 13, 2024 •

edited

Loading

markpollack Nov 13, 2024 •

edited

Loading

markpollack Nov 13, 2024 •

edited

Loading