Add optional token-budget embedding batches

vector/AGENTS.md

@@ -38,10 +38,15 @@ pipeline. Preserve these invariants when changing it.
 
 ## Fill batches without losing document boundaries
 
-- A positive `FillOptions.Batch.BatchSize` packs chunks across documents in
-  one scan page; values less than or equal to zero preserve the legacy
-  per-document encode unit. `BatchSize` remains the maximum texts in one
-  `EncodeFunc` call.
+- `WithFillBatch(WithBatchSize(n))` with a positive `n` packs chunks across
+  documents in one scan page. Omitting it preserves the legacy per-document
+  encode unit. `WithBatchSize` remains the maximum texts in one `EncodeFunc`
+  call.
+- `WithBatchTokenBudget` is opt-in and further reduces the effective batch
+  size from the caller's conservative per-input token upper bound. The vector
+  package does not choose a tokenizer, infer model limits, or alter input text.
+  Reject a configured upper bound that cannot fit one input before calling the
+  encoder.
 - Vectors from a shared encode batch must be scattered back to their exact
   document and chunk indexes before `SaveVectors`. Saves and `OnEncodeError`
   remain serialized and per document.

vector/encode.go

@@ -41,14 +41,50 @@ func (e *InvalidVectorError) Error() string {
 // backoff policy, since retryability is provider-specific.
 type EncodeFunc func(ctx context.Context, texts []string) ([][]float32, error)
 
-// BatchOptions controls how EncodeBatched groups and parallelizes calls.
-type BatchOptions struct {
-	// BatchSize is the maximum number of chunks passed to EncodeFunc in a
-	// single call. Values <= 0 send every chunk in one call.
-	BatchSize int
-	// Concurrency bounds how many EncodeFunc calls run at once. Values
-	// <= 0 mean one call at a time.
-	Concurrency int
+type batchOptions struct {
+	batchSize            int
+	concurrency          int
+	maxBatchTokens       int
+	inputTokenUpperBound int
+	tokenBudgetSet       bool
+}
+
+// BatchOption configures batching for EncodeBatched or Fill.
+type BatchOption func(*batchOptions)
+
+// WithBatchSize limits the number of chunks in one EncodeFunc call. Values
+// less than or equal to zero send every available chunk in one call.
+func WithBatchSize(size int) BatchOption {
+	return func(o *batchOptions) {
+		o.batchSize = size
+	}
+}
+
+// WithBatchConcurrency limits concurrent EncodeFunc calls. Values less than
+// or equal to zero use one call at a time.
+func WithBatchConcurrency(concurrency int) BatchOption {
+	return func(o *batchOptions) {
+		o.concurrency = concurrency
+	}
+}
+
+// WithBatchTokenBudget keeps an EncodeFunc call within an encoder's aggregate
+// input-token limit when every input is known to contain at most
+// inputTokenUpperBound tokens. Use it when the encoder enforces a combined
+// request limit that a count-only BatchSize cannot represent.
+//
+// The option caps the batch at maxBatchTokens / inputTokenUpperBound inputs,
+// or fewer when BatchSize is smaller. The caller must choose a conservative
+// per-input bound from its model and chunking rules. This package does not
+// count tokens, so a conservative bound can intentionally leave some request
+// capacity unused. Both values must be positive, and one input must fit within
+// maxBatchTokens.
+func WithBatchTokenBudget(maxBatchTokens, inputTokenUpperBound int) BatchOption {
+	return func(o *batchOptions) {
+		o.maxBatchTokens = maxBatchTokens
+		o.inputTokenUpperBound = inputTokenUpperBound
+		o.tokenBudgetSet = true
+	}
 }
 
 // EncodeBatched splits chunks into batches, invokes enc with bounded
@@ -61,10 +97,21 @@ type BatchOptions struct {
 // an error wrapping *InvalidVectorError, so faulty endpoint output never
 // reaches a Store. Blank chunk text is rejected with an error wrapping
 // ErrEmptyEmbeddingInput before any EncodeFunc call.
-func EncodeBatched(ctx context.Context, enc EncodeFunc, chunks []Chunk, o BatchOptions) ([]Vector, error) {
+func EncodeBatched(
+	ctx context.Context, enc EncodeFunc, chunks []Chunk, options ...BatchOption,
+) ([]Vector, error) {
+	return encodeBatched(ctx, enc, chunks, applyBatchOptions(options))
+}
+
+func encodeBatched(
+	ctx context.Context, enc EncodeFunc, chunks []Chunk, o batchOptions,
+) ([]Vector, error) {
 	if enc == nil {
 		return nil, fmt.Errorf("encode func is nil")
 	}
+	if err := o.validate(); err != nil {
+		return nil, err
+	}
 	if len(chunks) == 0 {
 		return nil, nil
 	}
@@ -74,11 +121,8 @@ func EncodeBatched(ctx context.Context, enc EncodeFunc, chunks []Chunk, o BatchO
 		}
 	}
 
-	batchSize := o.BatchSize
-	if batchSize <= 0 {
-		batchSize = len(chunks)
-	}
-	concurrency := o.Concurrency
+	batchSize := o.effectiveBatchSize(len(chunks))
+	concurrency := o.concurrency
 	if concurrency <= 0 {
 		concurrency = 1
 	}
@@ -168,6 +212,44 @@ launch:
 	return out, nil
 }
 
+func applyBatchOptions(options []BatchOption) batchOptions {
+	o := batchOptions{}
+	for _, option := range options {
+		if option != nil {
+			option(&o)
+		}
+	}
+	return o
+}
+
+func (o batchOptions) validate() error {
+	if !o.tokenBudgetSet {
+		return nil
+	}
+	if o.maxBatchTokens <= 0 || o.inputTokenUpperBound <= 0 {
+		return fmt.Errorf(
+			"batch token budget and input token upper bound must be positive: got %d and %d",
+			o.maxBatchTokens, o.inputTokenUpperBound)
+	}
+	if o.inputTokenUpperBound > o.maxBatchTokens {
+		return fmt.Errorf(
+			"input token upper bound %d exceeds batch token budget %d",
+			o.inputTokenUpperBound, o.maxBatchTokens)
+	}
+	return nil
+}
+
+func (o batchOptions) effectiveBatchSize(chunkCount int) int {
+	batchSize := o.batchSize
+	if batchSize <= 0 {
+		batchSize = chunkCount
+	}
+	if !o.tokenBudgetSet {
+		return batchSize
+	}
+	return min(batchSize, o.maxBatchTokens/o.inputTokenUpperBound)
+}
+
 // validateVector rejects vectors that would poison cosine distance: any
 // non-finite component, or a vector whose norm is zero. chunk is the global
 // chunk index reported in the error.

vector/encode_test.go

@@ -51,7 +51,8 @@ func TestEncodeBatchedPreservesOrderAcrossBatches(t *testing.T) {
 	})
 
 	in := chunks("a", "bb", "ccc", "dddd", "eeeee")
-	out, err := vector.EncodeBatched(context.Background(), enc, in, vector.BatchOptions{BatchSize: 2, Concurrency: 3})
+	out, err := vector.EncodeBatched(context.Background(), enc, in,
+		vector.WithBatchSize(2), vector.WithBatchConcurrency(3))
 	require.NoError(err)
 	require.Len(out, len(in))
 	for i, c := range in {
@@ -63,6 +64,39 @@ func TestEncodeBatchedPreservesOrderAcrossBatches(t *testing.T) {
 	assert.ElementsMatch([]int{2, 2, 1}, sizes, "batches are sized by BatchSize")
 }
 
+func TestEncodeBatchedCapsBatchSizeByTokenBudget(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	var batches [][]string
+	enc := echoEncoder(func(batch []string) {
+		batches = append(batches, append([]string(nil), batch...))
+	})
+	in := chunks("a", "bb", "ccc", "dddd")
+
+	out, err := vector.EncodeBatched(context.Background(), enc, in,
+		vector.WithBatchSize(4),
+		vector.WithBatchTokenBudget(120_000, 32_000),
+	)
+
+	require.NoError(err)
+	require.Len(out, len(in))
+	assert.Equal([][]string{{"a", "bb", "ccc"}, {"dddd"}}, batches)
+}
+
+func TestEncodeBatchedRejectsInputAboveTokenBudget(t *testing.T) {
+	var calls atomic.Int64
+	enc := echoEncoder(func([]string) { calls.Add(1) })
+
+	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a"),
+		vector.WithBatchTokenBudget(31_999, 32_000),
+	)
+
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "token budget")
+	assert.Zero(t, calls.Load(), "a known-oversized input is rejected before the provider call")
+}
+
 func TestEncodeBatchedRespectsConcurrencyBound(t *testing.T) {
 	require := require.New(t)
 	assert := assert.New(t)
@@ -85,7 +119,8 @@ func TestEncodeBatchedRespectsConcurrencyBound(t *testing.T) {
 	}
 
 	in := chunks("a", "b", "c", "d", "e", "f", "g", "h")
-	_, err := vector.EncodeBatched(context.Background(), enc, in, vector.BatchOptions{BatchSize: 1, Concurrency: 2})
+	_, err := vector.EncodeBatched(context.Background(), enc, in,
+		vector.WithBatchSize(1), vector.WithBatchConcurrency(2))
 	require.NoError(err)
 	assert.LessOrEqual(maxInFlight.Load(), int64(2), "never exceeds the concurrency bound")
 }
@@ -95,7 +130,7 @@ func TestEncodeBatchedSurfacesEncodeError(t *testing.T) {
 	sentinel := errors.New("boom")
 	enc := func(_ context.Context, _ []string) ([][]float32, error) { return nil, sentinel }
 
-	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b"), vector.BatchOptions{BatchSize: 1})
+	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b"), vector.WithBatchSize(1))
 	assert.ErrorIs(err, sentinel)
 }
 
@@ -119,10 +154,8 @@ func TestEncodeBatchedDoesNotLaunchBatchAfterBlockedDispatchSeesError(t *testing
 
 	done := make(chan error, 1)
 	go func() {
-		_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"), vector.BatchOptions{
-			BatchSize:   1,
-			Concurrency: 1,
-		})
+		_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"),
+			vector.WithBatchSize(1), vector.WithBatchConcurrency(1))
 		done <- err
 	}()
 
@@ -143,7 +176,7 @@ func TestEncodeBatchedRejectsCountMismatch(t *testing.T) {
 		return [][]float32{{1}}, nil // one vector for two texts
 	}
 
-	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b"), vector.BatchOptions{})
+	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b"))
 	assert.ErrorContains(err, "vectors for")
 }
 
@@ -163,7 +196,7 @@ func TestEncodeBatchedRejectsNonFiniteComponent(t *testing.T) {
 				return out, nil
 			}
 
-			_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"), vector.BatchOptions{})
+			_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"))
 			var invalid *vector.InvalidVectorError
 			require.ErrorAs(t, err, &invalid)
 			assert.Equal(t, 2, invalid.Chunk, "chunk index is global, not batch-relative")
@@ -187,20 +220,20 @@ func TestEncodeBatchedRejectsZeroNormVector(t *testing.T) {
 
 	// BatchSize 2 puts the zero vector in the second batch, so a
 	// batch-relative index would wrongly report 0.
-	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"), vector.BatchOptions{BatchSize: 2})
+	_, err := vector.EncodeBatched(context.Background(), enc, chunks("a", "b", "c"), vector.WithBatchSize(2))
 	var invalid *vector.InvalidVectorError
 	require.ErrorAs(t, err, &invalid)
 	assert.Equal(t, 2, invalid.Chunk)
 	assert.Equal(t, -1, invalid.Component, "zero norm reports no single component")
 }
 
 func TestEncodeBatchedNilEncoder(t *testing.T) {
-	_, err := vector.EncodeBatched(context.Background(), nil, chunks("a"), vector.BatchOptions{})
+	_, err := vector.EncodeBatched(context.Background(), nil, chunks("a"))
 	assert.Error(t, err)
 }
 
 func TestEncodeBatchedEmptyInput(t *testing.T) {
-	out, err := vector.EncodeBatched(context.Background(), echoEncoder(nil), nil, vector.BatchOptions{})
+	out, err := vector.EncodeBatched(context.Background(), echoEncoder(nil), nil)
 	require.NoError(t, err)
 	assert.Empty(t, out)
 }
@@ -209,10 +242,8 @@ func TestEncodeBatchedRejectsWhitespaceBeforeCallingEncoder(t *testing.T) {
 	var calls atomic.Int64
 	enc := echoEncoder(func([]string) { calls.Add(1) })
 
-	_, err := vector.EncodeBatched(context.Background(), enc, chunks("alpha", " \t\n\u2003"), vector.BatchOptions{
-		BatchSize:   1,
-		Concurrency: 2,
-	})
+	_, err := vector.EncodeBatched(context.Background(), enc, chunks("alpha", " \t\n\u2003"),
+		vector.WithBatchSize(1), vector.WithBatchConcurrency(2))
 
 	require.ErrorIs(t, err, vector.ErrEmptyEmbeddingInput)
 	assert.Zero(t, calls.Load(), "the batch is validated before any provider call")
@@ -227,10 +258,8 @@ func TestEncodeBatchedRejectsInvisibleTextBeforeCallingEncoder(t *testing.T) {
 	enc := echoEncoder(func([]string) { calls.Add(1) })
 
 	_, err := vector.EncodeBatched(context.Background(), enc,
-		chunks("alpha", "\u200b\ufeff\u200d"), vector.BatchOptions{
-			BatchSize:   1,
-			Concurrency: 2,
-		})
+		chunks("alpha", "\u200b\ufeff\u200d"),
+		vector.WithBatchSize(1), vector.WithBatchConcurrency(2))
 
 	require.ErrorIs(t, err, vector.ErrEmptyEmbeddingInput)
 	assert.Zero(t, calls.Load(), "the batch is validated before any provider call")
@@ -239,6 +268,6 @@ func TestEncodeBatchedRejectsInvisibleTextBeforeCallingEncoder(t *testing.T) {
 func TestEncodeBatchedStopsOnCancelledContext(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
-	_, err := vector.EncodeBatched(ctx, echoEncoder(nil), chunks("a", "b"), vector.BatchOptions{BatchSize: 1})
+	_, err := vector.EncodeBatched(ctx, echoEncoder(nil), chunks("a", "b"), vector.WithBatchSize(1))
 	assert.ErrorIs(t, err, context.Canceled)
 }