From 33ac9224bfe613da91fa46479a86a214cd09e0b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20K=C5=99=C3=ADstek?=
 <77566262+David-Kristek@users.noreply.github.com>
Date: Fri, 26 Apr 2024 19:57:46 +0200
Subject: [PATCH] fix(llama-index): avoid batching in embeddings (#353)

---
 src/genai/extensions/llama_index/embeddings.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/genai/extensions/llama_index/embeddings.py b/src/genai/extensions/llama_index/embeddings.py
index 195829fd..3cb7dd85 100644
--- a/src/genai/extensions/llama_index/embeddings.py
+++ b/src/genai/extensions/llama_index/embeddings.py
@@ -1,6 +1,8 @@
 import asyncio
 from typing import Optional
 
+from pydantic import Field
+
 from genai._types import ModelLike
 from genai.client import Client
 from genai.schema import TextEmbeddingParameters
@@ -17,6 +19,9 @@ class IBMGenAILlamaIndexEmbedding(BaseEmbedding):
     model_id: str
     parameters: Optional[ModelLike[TextEmbeddingParameters]] = None
     execution_options: Optional[ModelLike[CreateExecutionOptions]] = None
+    # Batch size is set to 100000 to avoid batching in
+    # LlamaIndex as it is handled by the SDK itself
+    embed_batch_size: int = Field(default=10000, description="The batch size for embedding calls.", gt=0)
 
     @classmethod
     def class_name(cls) -> str: