Easy-rag: allow to specify minScore through configuration

quarkiverse · Oct 31, 2024 · 99297ed · 99297ed
1 parent fd98b21
commit 99297ed
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/...y-rag/runtime/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRagConfig.java b/...y-rag/runtime/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRagConfig.java
@@ -2,7 +2,8 @@
 
 import static io.quarkus.runtime.annotations.ConfigPhase.RUN_TIME;
 
-import io.quarkiverse.langchain4j.easyrag.EasyRagManualIngestion;
+import java.util.OptionalDouble;
+
 import io.quarkus.runtime.annotations.ConfigGroup;
 import io.quarkus.runtime.annotations.ConfigRoot;
 import io.smallrye.config.ConfigMapping;
@@ -52,6 +53,11 @@ public interface EasyRagConfig {
     @WithDefault("5")
     Integer maxResults();
 
+    /**
+     * The minimum score for results to return when querying the retrieval augmentor.
+     */
+    OptionalDouble minScore();
+
     /**
      * The strategy to decide whether document ingestion into the store
      * should happen at startup or not. The default is ON. Changing to OFF

diff --git a/...rag/runtime/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRagRecorder.java b/...rag/runtime/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRagRecorder.java
@@ -72,7 +72,7 @@ public RetrievalAugmentor apply(SyntheticCreationalContext<RetrievalAugmentor> c
                 EmbeddingModel model = context.getInjectedReference(EmbeddingModel.class, Default.Literal.INSTANCE);
                 EmbeddingStore<TextSegment> store = context.getInjectedReference(EmbeddingStore.class,
                         Default.Literal.INSTANCE);
-                return new EasyRetrievalAugmentor(config.maxResults(), model, store);
+                return new EasyRetrievalAugmentor(config, model, store);
             }
         };
     }

diff --git a/...time/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRetrievalAugmentor.java b/...time/src/main/java/io/quarkiverse/langchain4j/easyrag/runtime/EasyRetrievalAugmentor.java
@@ -16,14 +16,16 @@ public class EasyRetrievalAugmentor implements RetrievalAugmentor {
 
     private DefaultRetrievalAugmentor delegate;
 
-    public EasyRetrievalAugmentor(Integer maxResults, EmbeddingModel embeddingModel, EmbeddingStore embeddingStore) {
-        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
+    public EasyRetrievalAugmentor(EasyRagConfig config, EmbeddingModel embeddingModel, EmbeddingStore embeddingStore) {
+        var contentRetrieverBuilder = EmbeddingStoreContentRetriever.builder()
                 .embeddingModel(embeddingModel)
                 .embeddingStore(embeddingStore)
-                .maxResults(maxResults)
-                .build();
+                .maxResults(config.maxResults());
+
+        config.minScore().ifPresent(contentRetrieverBuilder::minScore);
+
         delegate = DefaultRetrievalAugmentor.builder()
-                .contentRetriever(contentRetriever).build();
+                .contentRetriever(contentRetrieverBuilder.build()).build();
     }
 
     @Override