Add RAG test using QueryRouter and ContentInjector

Author: mariofusco

integration-tests/rag/src/main/java/org/acme/example/AiServiceWithQueryRouterAndContentInjector.java

@@ -0,0 +1,78 @@
+package org.acme.example;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Supplier;
+
+import jakarta.inject.Singleton;
+
+import dev.langchain4j.data.message.ChatMessage;
+import dev.langchain4j.data.message.UserMessage;
+import dev.langchain4j.rag.DefaultRetrievalAugmentor;
+import dev.langchain4j.rag.RetrievalAugmentor;
+import dev.langchain4j.rag.content.Content;
+import dev.langchain4j.rag.content.injector.ContentInjector;
+import dev.langchain4j.rag.content.retriever.ContentRetriever;
+import dev.langchain4j.rag.query.Query;
+import dev.langchain4j.rag.query.router.QueryRouter;
+import io.quarkiverse.langchain4j.RegisterAiService;
+
+@RegisterAiService(retrievalAugmentor = AiServiceWithQueryRouterAndContentInjector.QueryRouterAugmentor.class)
+public interface AiServiceWithQueryRouterAndContentInjector {
+
+    String chat(String message);
+
+    /**
+     * Contains a query transformer that transforms the query text to
+     * lowercase. If the transformed worked properly, the content retriever
+     * will append content saying "The transformer works!".
+     */
+    @Singleton
+    class QueryRouterAugmentor implements Supplier<RetrievalAugmentor> {
+
+        @Override
+        public RetrievalAugmentor get() {
+            return DefaultRetrievalAugmentor.builder()
+                    .queryRouter(new QueryRouter() {
+                        @Override
+                        public Collection<ContentRetriever> route(Query query) {
+                            if (query.text().contains("dog")) {
+                                return Collections.singletonList(dogsRetriever());
+                            } else if (query.text().contains("cat")) {
+                                return Collections.singletonList(catsRetriever());
+                            } else {
+                                return Collections.emptyList();
+                            }
+                        }
+                    })
+                    .contentInjector(new ContentInjector() {
+                        @Override
+                        public ChatMessage inject(List<Content> contents, ChatMessage chatMessage) {
+                            String rewrittenMessage = ((UserMessage) chatMessage).singleText() + " - "
+                                    + contents.get(0).textSegment().text();
+                            return UserMessage.userMessage(rewrittenMessage);
+                        }
+                    })
+                    .build();
+        }
+
+        static ContentRetriever dogsRetriever() {
+            return new ContentRetriever() {
+                @Override
+                public List<Content> retrieve(Query query) {
+                    return Collections.singletonList(Content.from("Dogs bark"));
+                }
+            };
+        }
+
+        static ContentRetriever catsRetriever() {
+            return new ContentRetriever() {
+                @Override
+                public List<Content> retrieve(Query query) {
+                    return Collections.singletonList(Content.from("Cats meow"));
+                }
+            };
+        }
+    }
+}

integration-tests/rag/src/test/java/org/acme/example/openai/RAGWithQueryRouterAndContentInjectorTest.java

@@ -0,0 +1,51 @@
+package org.acme.example.openai;
+
+import static io.quarkiverse.langchain4j.runtime.LangChain4jUtil.chatMessageToText;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import jakarta.inject.Inject;
+
+import org.acme.example.AiServiceWithQueryRouterAndContentInjector;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.mockito.stubbing.Answer;
+
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.data.message.ChatMessage;
+import dev.langchain4j.model.chat.ChatModel;
+import dev.langchain4j.model.chat.request.ChatRequest;
+import dev.langchain4j.model.chat.response.ChatResponse;
+import io.quarkus.test.junit.QuarkusMock;
+import io.quarkus.test.junit.QuarkusTest;
+
+@QuarkusTest
+public class RAGWithQueryRouterAndContentInjectorTest {
+
+    @Inject
+    AiServiceWithQueryRouterAndContentInjector service;
+
+    private static AtomicReference<List<ChatMessage>> lastQuery = new AtomicReference<>();
+
+    @BeforeAll
+    public static void initializeModel() {
+        ChatModel mock = Mockito.mock(ChatModel.class);
+        Answer<ChatResponse> answer = invocation -> {
+            lastQuery.set(((ChatRequest) invocation.getArgument(0)).messages());
+            return ChatResponse.builder().aiMessage(new AiMessage("Mock response")).build();
+        };
+        Mockito.when(mock.chat(Mockito.any(ChatRequest.class))).thenAnswer(answer);
+        QuarkusMock.installMockForType(mock, ChatModel.class);
+    }
+
+    @Test
+    public void test() {
+        service.chat("What dogs do?");
+        String query = chatMessageToText(lastQuery.get().get(0));
+        Assertions.assertTrue(query.equals("What dogs do? - Dogs bark"), query);
+    }
+
+}