From 0825d52224dfb3091c162d9d266c4321dcc9e9f4 Mon Sep 17 00:00:00 2001
From: 1993heqiang <531364804@qq.com>
Date: Tue, 5 Nov 2024 00:10:17 +0800
Subject: [PATCH] Add validation and rename parameter in
 TokenCountBatchingStrategy

This commit rename the thresholdFactor parameter to reservePercentage to better
reflect its purpose in the TokenCountBatchingStrategy class. It also adds
validation for input parameters including maxInputTokenCount > 0 and
reservePercentage between 0 and 1 for safer initialization.

The change ensures proper parameter validation on TokenCountBatchingStrategy
creation to prevent potential runtime errors from invalid inputs.
---
 .../ai/embedding/TokenCountBatchingStrategy.java  | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
index 298278c35..8a2a1195c 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
@@ -78,12 +78,13 @@ public class TokenCountBatchingStrategy implements BatchingStrategy {
 
 	/**
 	 * @param encodingType {@link EncodingType}
-	 * @param thresholdFactor the threshold factor to use on top of the max input token
-	 * count
+	 * @param reservePercentage the percentage of tokens to reserve from the max input
+	 * token count to create a buffer.
 	 * @param maxInputTokenCount upper limit for input tokens
 	 */
-	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double thresholdFactor) {
-		this(encodingType, maxInputTokenCount, thresholdFactor, Document.DEFAULT_CONTENT_FORMATTER, MetadataMode.NONE);
+	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage) {
+		this(encodingType, maxInputTokenCount, reservePercentage, Document.DEFAULT_CONTENT_FORMATTER,
+				MetadataMode.NONE);
 	}
 
 	/**
@@ -99,6 +100,8 @@ public class TokenCountBatchingStrategy implements BatchingStrategy {
 	public TokenCountBatchingStrategy(EncodingType encodingType, int maxInputTokenCount, double reservePercentage,
 			ContentFormatter contentFormatter, MetadataMode metadataMode) {
 		Assert.notNull(encodingType, "EncodingType must not be null");
+		Assert.isTrue(maxInputTokenCount > 0, "MaxInputTokenCount must be greater than 0");
+		Assert.isTrue(reservePercentage >= 0 && reservePercentage < 1, "ReservePercentage must be in range [0, 1)");
 		Assert.notNull(contentFormatter, "ContentFormatter must not be null");
 		Assert.notNull(metadataMode, "MetadataMode must not be null");
 		this.tokenCountEstimator = new JTokkitTokenCountEstimator(encodingType);
@@ -120,6 +123,10 @@ public class TokenCountBatchingStrategy implements BatchingStrategy {
 	public TokenCountBatchingStrategy(TokenCountEstimator tokenCountEstimator, int maxInputTokenCount,
 			double reservePercentage, ContentFormatter contentFormatter, MetadataMode metadataMode) {
 		Assert.notNull(tokenCountEstimator, "TokenCountEstimator must not be null");
+		Assert.isTrue(maxInputTokenCount > 0, "MaxInputTokenCount must be greater than 0");
+		Assert.isTrue(reservePercentage >= 0 && reservePercentage < 1, "ReservePercentage must be in range [0, 1)");
+		Assert.notNull(contentFormatter, "ContentFormatter must not be null");
+		Assert.notNull(metadataMode, "MetadataMode must not be null");
 		this.tokenCountEstimator = tokenCountEstimator;
 		this.maxInputTokenCount = (int) Math.round(maxInputTokenCount * (1 - reservePercentage));
 		this.contentFormater = contentFormatter;