fix(wren-ai-service): Optimize SQL Data Preprocessing with Progressiv…

…e Data Reduction (#1331)

wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py

@@ -19,23 +19,51 @@ def preprocess(
     sql_data: Dict,
     encoding: tiktoken.Encoding,
 ) -> Dict:
+    def reduce_data_size(data: list, reduction_step: int = 50) -> list:
+        """Reduce the size of data by removing elements from the end.
+
+        Args:
+            data: The input list to reduce
+            reduction_step: Number of elements to remove (must be positive)
+
+        Returns:
+            list: A list with reduced size
+
+        Raises:
+            ValueError: If reduction_step is not positive
+        """
+        if reduction_step <= 0:
+            raise ValueError("reduction_step must be positive")
+
+        elements_to_keep = max(0, len(data) - reduction_step)
+        returned_data = data[:elements_to_keep]
+
+        logger.info(
+            f"Reducing data size by {reduction_step} rows. "
+            f"Original size: {len(data)}, New size: {len(returned_data)}"
+        )
+
+        return returned_data
+
     _token_count = len(encoding.encode(str(sql_data)))
     num_rows_used_in_llm = len(sql_data.get("data", []))
+    iteration = 0
 
-    if _token_count > 100_000:
-        sql_data = {
-            "columns": sql_data.get("columns", []),
-            "data": sql_data.get("data", [])[:250],
-            "dtypes": sql_data.get("dtypes", {}),
-        }
+    while _token_count > 100_000:
+        if iteration > 1000:
+            """
+            Avoid infinite loop
+            If the token count is still too high after 1000 iterations, break
+            """
+            break
 
-        num_rows_used_in_llm = len(sql_data.get("data", []))
+        iteration += 1
 
-        return {
-            "sql_data": sql_data,
-            "num_rows_used_in_llm": num_rows_used_in_llm,
-            "tokens": _token_count,
-        }
+        data = sql_data.get("data", [])
+        sql_data["data"] = reduce_data_size(data)
+        num_rows_used_in_llm = len(sql_data.get("data", []))
+        _token_count = len(encoding.encode(str(sql_data)))
+        logger.info(f"Token count: {_token_count}")
 
     return {
         "sql_data": sql_data,