KB max batch size: apply insert limit globally (#11905)

Author: ea-rus

mindsdb/interfaces/agents/constants.py

@@ -208,7 +208,7 @@
 USER_COLUMN = "question"
 DEFAULT_EMBEDDINGS_MODEL_PROVIDER = "openai"
 DEFAULT_EMBEDDINGS_MODEL_CLASS = OpenAIEmbeddings
-MAX_INSERT_BATCH_SIZE = 50_000
+MAX_INSERT_BATCH_SIZE = int(os.getenv("KB_MAX_INSERT_BATCH_SIZE", 50_000))
 DEFAULT_TIKTOKEN_MODEL_NAME = os.getenv("DEFAULT_TIKTOKEN_MODEL_NAME", "gpt-4")
 AGENT_CHUNK_POLLING_INTERVAL_SECONDS = os.getenv("AGENT_CHUNK_POLLING_INTERVAL_SECONDS", 1.0)
 DEFAULT_TEXT2SQL_DATABASE = "mindsdb"

mindsdb/interfaces/knowledge_base/controller.py

@@ -586,8 +586,6 @@ def insert_rows(self, rows: List[Dict]):
         """Process and insert raw data rows"""
         if not rows:
             return
-        if len(rows) > MAX_INSERT_BATCH_SIZE:
-            raise ValueError("Input data is too large, please load data in batches")
 
         df = pd.DataFrame(rows)
 
@@ -700,6 +698,9 @@ def insert(self, df: pd.DataFrame, params: dict = None):
         if df.empty:
             return
 
+        if len(df) > MAX_INSERT_BATCH_SIZE:
+            raise ValueError("Input data is too large, please load data in batches")
+
         try:
             run_query_id = ctx.run_query_id
             # Link current KB to running query (where KB is used to insert data)