From b1cf4b7dd3c61856bdeab0336bf9844aa55d9a62 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B8=A3=E5=B1=B1?= <dengjiaji.djj@antgroup.com>
Date: Mon, 16 Jun 2025 15:16:44 +0800
Subject: [PATCH] fix bugs for large batch size of embeddings

---
 experiencemaker/model/base_embedding_model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/experiencemaker/model/base_embedding_model.py b/experiencemaker/model/base_embedding_model.py
index c611d11b..93891900 100644
--- a/experiencemaker/model/base_embedding_model.py
+++ b/experiencemaker/model/base_embedding_model.py
@@ -76,7 +76,9 @@ class BaseEmbeddingModel(BaseModel, ABC):
             return nodes
 
         elif isinstance(nodes, list):
-            embeddings = self.get_embeddings(input_text=[node.content for node in nodes])
+            max_batch_size = 10 # text-embedding-v4 batch size should not be larger than 10
+            embeddings = [emb for i in range(0, len(nodes), max_batch_size) for emb in
+                          self.get_embeddings(input_text=[node.content for node in nodes[i:i + max_batch_size]])]
             if len(embeddings) != len(nodes):
                 logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}")
             else: