From 44e8a8cf536aa3a23b2c3fd1a962f33f1f8f889e Mon Sep 17 00:00:00 2001
From: Jan Martiska <jmartisk@redhat.com>
Date: Thu, 15 Feb 2024 13:12:51 +0100
Subject: [PATCH] Use batching in Redis to avoid overflowing the connection
 pool

---
 .../langchain4j/redis/RedisEmbeddingStore.java           | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/redis/runtime/src/main/java/io/quarkiverse/langchain4j/redis/RedisEmbeddingStore.java b/redis/runtime/src/main/java/io/quarkiverse/langchain4j/redis/RedisEmbeddingStore.java
index 6161dd202..bd80942ae 100644
--- a/redis/runtime/src/main/java/io/quarkiverse/langchain4j/redis/RedisEmbeddingStore.java
+++ b/redis/runtime/src/main/java/io/quarkiverse/langchain4j/redis/RedisEmbeddingStore.java
@@ -5,6 +5,7 @@
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -25,7 +26,6 @@
 import io.quarkiverse.langchain4j.QuarkusJsonCodecFactory;
 import io.quarkiverse.langchain4j.redis.runtime.RedisSchema;
 import io.quarkus.redis.datasource.ReactiveRedisDataSource;
-import io.quarkus.redis.datasource.json.ReactiveJsonCommands;
 import io.quarkus.redis.datasource.keys.KeyScanArgs;
 import io.quarkus.redis.datasource.search.CreateArgs;
 import io.quarkus.redis.datasource.search.Document;
@@ -125,9 +125,8 @@ private void addAllInternal(List<String> ids, List<Embedding> embeddings, List<T
         if (ids.isEmpty() || ids.size() != embeddings.size() || (embedded != null && embedded.size() != embeddings.size())) {
             throw new IllegalArgumentException("ids, embeddings and embedded must be non-empty and of the same size");
         }
-        ReactiveJsonCommands<String> json = ds.json();
         int size = ids.size();
-        Uni[] unis = new Uni[size];
+        List<Request> commands = new ArrayList<>();
         for (int i = 0; i < size; i++) {
             String id = ids.get(i);
             Embedding embedding = embeddings.get(i);
@@ -147,9 +146,9 @@ private void addAllInternal(List<String> ids, List<Embedding> embeddings, List<T
                 fields.putAll(textSegment.metadata().asMap());
             }
             String key = schema.getPrefix() + id;
-            unis[i] = json.jsonSet(key, "$", fields);
+            commands.add(Request.cmd(Command.JSON_SET).arg(key).arg("$").arg(Json.toJson(fields)));
         }
-        Uni.join().all(unis).andFailFast().await().indefinitely();
+        ds.getRedis().batchAndAwait(commands);
     }
 
     @Override