Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add validation and rename parameter in TokenCountBatchingStrategy

Browse Source 
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.
This commit is contained in:
1993heqiang
2024-11-05 00:10:17 +08:00
committed by Mark Pollack
parent 83f7164bff
commit 0825d52224
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
spring-ai-core/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java
View File

@@ -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;