From d80fd64a1876e089120859222a9fddc670730382 Mon Sep 17 00:00:00 2001 From: John Blum Date: Sun, 9 Nov 2025 17:02:28 -0800 Subject: [PATCH] Provide detailed Exception message when token count exceeds max Closes #4835 Signed-off-by: John Blum --- .../MaxTokenCountExceededException.java | 54 +++++++++++++++++++ .../embedding/TokenCountBatchingStrategy.java | 4 +- .../TokenCountBatchingStrategyTests.java | 41 +++++++++++++- 3 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 spring-ai-commons/src/main/java/org/springframework/ai/embedding/MaxTokenCountExceededException.java diff --git a/spring-ai-commons/src/main/java/org/springframework/ai/embedding/MaxTokenCountExceededException.java b/spring-ai-commons/src/main/java/org/springframework/ai/embedding/MaxTokenCountExceededException.java new file mode 100644 index 00000000000..533fa2f9df3 --- /dev/null +++ b/spring-ai-commons/src/main/java/org/springframework/ai/embedding/MaxTokenCountExceededException.java @@ -0,0 +1,54 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.embedding; + +import org.springframework.ai.document.Document; + +/** + * {@link RuntimeException} thrown when the token count of the provided content exceeds + * the configured maximum. + * + * @author John Blum + * @see IllegalArgumentException + * @since 1.1.0 + */ +@SuppressWarnings("unused") +public class MaxTokenCountExceededException extends IllegalArgumentException { + + public static MaxTokenCountExceededException because(Document document, int tokenCount, int maxTokenCount) { + String message = "Tokens [%d] from Document [%s] exceeds the configured maximum number of input tokens allowed [%d]" + .formatted(tokenCount, document.getId(), maxTokenCount); + return new MaxTokenCountExceededException(message); + } + + public MaxTokenCountExceededException() { + + } + + public MaxTokenCountExceededException(String message) { + super(message); + } + + public MaxTokenCountExceededException(Throwable cause) { + super(cause); + } + + public MaxTokenCountExceededException(String message, Throwable cause) { + super(message, cause); + } + +} diff --git a/spring-ai-commons/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java b/spring-ai-commons/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java index f421b11591e..3e30172bf06 100644 --- a/spring-ai-commons/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java +++ b/spring-ai-commons/src/main/java/org/springframework/ai/embedding/TokenCountBatchingStrategy.java @@ -51,6 +51,7 @@ * @author Laura Trotta * @author Jihoon Kim * @author Yanming Zhou + * @author John Blum * @since 1.0.0 */ public class TokenCountBatchingStrategy implements BatchingStrategy { @@ -148,8 +149,7 @@ public List> batch(List documents) { int tokenCount = this.tokenCountEstimator .estimate(document.getFormattedContent(this.contentFormatter, this.metadataMode)); if (tokenCount > this.maxInputTokenCount) { - throw new IllegalArgumentException( - "Tokens in a single document exceeds the maximum number of allowed input tokens"); + throw MaxTokenCountExceededException.because(document, tokenCount, this.maxInputTokenCount); } documentTokens.put(document, tokenCount); } diff --git a/spring-ai-commons/src/test/java/org/springframework/ai/embedding/TokenCountBatchingStrategyTests.java b/spring-ai-commons/src/test/java/org/springframework/ai/embedding/TokenCountBatchingStrategyTests.java index 3d14f1e7125..4a0e2563571 100644 --- a/spring-ai-commons/src/test/java/org/springframework/ai/embedding/TokenCountBatchingStrategyTests.java +++ b/spring-ai-commons/src/test/java/org/springframework/ai/embedding/TokenCountBatchingStrategyTests.java @@ -22,17 +22,30 @@ import org.junit.jupiter.api.Test; +import org.springframework.ai.document.ContentFormatter; import org.springframework.ai.document.Document; +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.tokenizer.TokenCountEstimator; import org.springframework.core.io.DefaultResourceLoader; import org.springframework.core.io.Resource; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; /** - * Basic unit test for {@link TokenCountBatchingStrategy}. + * Basic unit tests for {@link TokenCountBatchingStrategy}. * * @author Soby Chacko + * @author John Blum */ public class TokenCountBatchingStrategyTests { @@ -54,4 +67,30 @@ void batchEmbeddingWithLargeDocumentExceedsMaxTokenSize() throws IOException { .isInstanceOf(IllegalArgumentException.class); } + @Test + void documentTokenCountExceedsConfiguredMaxTokenCount() { + + Document mockDocument = mock(Document.class); + ContentFormatter mockContentFormatter = mock(ContentFormatter.class); + TokenCountEstimator mockTokenCountEstimator = mock(TokenCountEstimator.class); + + doReturn("123abc").when(mockDocument).getId(); + doReturn(10).when(mockTokenCountEstimator).estimate(anyString()); + doReturn("test").when(mockDocument).getFormattedContent(any(), any()); + + TokenCountBatchingStrategy batchingStrategy = new TokenCountBatchingStrategy(mockTokenCountEstimator, 9, 0.0d, + mockContentFormatter, MetadataMode.EMBED); + + assertThatExceptionOfType(MaxTokenCountExceededException.class) + .isThrownBy(() -> batchingStrategy.batch(List.of(mockDocument))) + .withMessage( + "Tokens [10] from Document [123abc] exceeds the configured maximum number of input tokens allowed [9]") + .withNoCause(); + + verify(mockDocument, times(1)).getId(); + verify(mockDocument, times(1)).getFormattedContent(eq(mockContentFormatter), eq(MetadataMode.EMBED)); + verify(mockTokenCountEstimator, times(1)).estimate(eq("test")); + verifyNoMoreInteractions(mockDocument, mockTokenCountEstimator); + } + }