Skip to content

Fix TokenCountBatchingStrategy ignoring metadata when sizing batches - #6782

Open
subhashpolisetti wants to merge 1 commit into
spring-projects:mainfrom
subhashpolisetti:fix-batching-strategy-metadata-mode-default
Open

Fix TokenCountBatchingStrategy ignoring metadata when sizing batches#6782
subhashpolisetti wants to merge 1 commit into
spring-projects:mainfrom
subhashpolisetti:fix-batching-strategy-metadata-mode-default

Conversation

@subhashpolisetti

Copy link
Copy Markdown
Contributor

The metadata-aware embedding models send
Document.getFormattedContent(MetadataMode.EMBED) to the embedding API, but
TokenCountBatchingStrategy estimated token counts with MetadataMode.NONE. Both are
auto-configured defaults in the same application context: OpenAiEmbeddingProperties binds
MetadataMode.EMBED into the model, while every vector store auto-configuration registers
the no-argument TokenCountBatchingStrategy. The default path therefore counted the document
text and then submitted that text plus its metadata.

Batches exceeded the limit they were sized for, and the guard that rejects an oversized
document measured a different string than the one submitted: with the default 8191 token
count and 10% reserve, a document counted at 7001 tokens is accepted and then sent as 8440,
past both the 7371 token budget the batch was sized for and the model's own input limit.

Estimate with MetadataMode.EMBED instead. For models that do not send metadata the estimate
becomes conservative rather than wrong, so batches stay within the limit either way. The
constructor taking an explicit MetadataMode is unchanged, so an application embedding text
only can still request MetadataMode.NONE directly.

Testing

Added batchRejectsDocumentWhoseMetadataPushesItOverTheLimit and
batchSplitsOnTheTokenCountIncludingMetadata, which both fail without the change (the metadata
is not counted, so the oversized document is accepted and the three documents are placed in a
single batch) and pass with it. batchWithExplicitMetadataModeNoneIgnoresMetadata covers the
opt-out and passes either way, pinning the behaviour of the constructor that takes an explicit
MetadataMode. Existing batching tests are unaffected. No credentials are needed: the tests use
the real JTokkitTokenCountEstimator with small token limits.

./mvnw -pl spring-ai-model clean test passes (837 tests).

The metadata-aware embedding models send
`Document.getFormattedContent(MetadataMode.EMBED)` to the embedding API,
but `TokenCountBatchingStrategy` estimated token counts with
`MetadataMode.NONE`. Both are auto-configured defaults in the same
application context: `OpenAiEmbeddingProperties` binds `MetadataMode.EMBED`
into the model, while every vector store auto-configuration registers the
no-argument `TokenCountBatchingStrategy`. The default path therefore
counted the document text and then submitted that text plus its metadata.

Batches exceeded the limit they were sized for, and the guard that rejects
an oversized document measured a different string than the one submitted:
with the default 8191 token count and 10% reserve, a document counted at
7001 tokens is accepted and then sent as 8440, past both the 7371 token
budget the batch was sized for and the model's own input limit.

Estimate with `MetadataMode.EMBED` instead. For models that do not send
metadata the estimate becomes conservative rather than wrong, so batches
stay within the limit either way. The constructor taking an explicit
`MetadataMode` is unchanged, so an application embedding text only can
still request `MetadataMode.NONE` directly.

Signed-off-by: subhash polisetti <subhashr161347@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants