From cf0278ca43ab3bf2531195c116b13543c852e3f6 Mon Sep 17 00:00:00 2001
From: billvsme <994171686@qq.com>
Date: Fri, 13 Dec 2024 15:42:41 +0800
Subject: [PATCH] The asyncio.as_completed() function does not guarantee that
 the results are ordered

---
 lightrag/kg/milvus_impl.py | 2 +-
 lightrag/storage.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/lightrag/kg/milvus_impl.py b/lightrag/kg/milvus_impl.py
index 6d2520ce..c72498c5 100644
--- a/lightrag/kg/milvus_impl.py
+++ b/lightrag/kg/milvus_impl.py
@@ -59,7 +59,7 @@ class MilvusVectorDBStorge(BaseVectorStorage):
         embedding_tasks = [self.embedding_func(batch) for batch in batches]
         embeddings_list = []
         for f in tqdm_async(
-            asyncio.as_completed(embedding_tasks),
+            await asyncio.gather(*embedding_tasks),
             total=len(embedding_tasks),
             desc="Generating embeddings",
             unit="batch",
diff --git a/lightrag/storage.py b/lightrag/storage.py
index 4c043893..534c6e2e 100644
--- a/lightrag/storage.py
+++ b/lightrag/storage.py
@@ -99,7 +99,7 @@ class NanoVectorDBStorage(BaseVectorStorage):
         embedding_tasks = [self.embedding_func(batch) for batch in batches]
         embeddings_list = []
         for f in tqdm_async(
-            asyncio.as_completed(embedding_tasks),
+            await asyncio.gather(*embedding_tasks),
             total=len(embedding_tasks),
             desc="Generating embeddings",
             unit="batch",