diff --git a/vector/AGENTS.md b/vector/AGENTS.md index 17ce2f9..bfea0dd 100644 --- a/vector/AGENTS.md +++ b/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. diff --git a/vector/encode.go b/vector/encode.go index c1bb64c..4d57eab 100644 --- a/vector/encode.go +++ b/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. diff --git a/vector/encode_test.go b/vector/encode_test.go index 83c37e2..f1e6ecc 100644 --- a/vector/encode_test.go +++ b/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,7 +220,7 @@ 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) @@ -195,12 +228,12 @@ func TestEncodeBatchedRejectsZeroNormVector(t *testing.T) { } 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) } diff --git a/vector/flow.go b/vector/flow.go index 0ab7f3c..59cceeb 100644 --- a/vector/flow.go +++ b/vector/flow.go @@ -7,44 +7,70 @@ import ( "sync" ) -// FillOptions configures Fill. K is the document key type of the Store the -// fill runs over. -type FillOptions[K comparable] struct { - // ScanBatch is the number of pending documents fetched per scan. - // Values <= 0 use 128. - ScanBatch int - // Split controls how each document's content is windowed into chunks. - Split SplitOptions - // Batch controls how chunks are batched into encode calls. A positive - // BatchSize packs chunks across documents within each scan page. Values - // <= 0 preserve the per-document encode unit. - Batch BatchOptions - // Concurrency bounds parallel fill work within each scan page. Values - // <= 0 use 1 (sequential). - // SaveVectors and OnEncodeError stay serialized on the calling goroutine - // regardless, so stores and hooks need no extra locking. With a positive - // BatchSize, it composes with Batch.Concurrency to bound concurrent encode - // calls by their product. - Concurrency int - // OnEncodeError, if non-nil, is consulted when encoding a document - // fails. Returning true skips the document: it is stamped for the - // generation with no vectors so it stops being pending, and the fill - // continues (the treatment for inputs a model permanently rejects). - // Returning false — or leaving OnEncodeError nil — aborts the fill - // with the error, which is the right default for transient failures. - OnEncodeError func(doc K, err error) bool - // ShouldIsolateBatchError reports whether an error from a shared, - // multi-document encode call might be caused by an individual input and - // is worth diagnosing at document-slice granularity. - // - // Fill does not call this function for single-document batches or for - // errors wrapping context.Canceled or context.DeadlineExceeded. Errors - // carrying exact batch-position attribution (*InvalidVectorError) also - // bypass this function and go directly to OnEncodeError. Returning true - // permits diagnosis only; OnEncodeError still owns the decision to - // skip-stamp an attributed document. Returning false, or leaving this nil, - // aborts Fill without document-level retries. - ShouldIsolateBatchError func(error) bool +type fillOptions[K comparable] struct { + scanBatch int + split SplitOptions + batch batchOptions + concurrency int + onEncodeError func(doc K, err error) bool + shouldIsolateBatchError func(error) bool +} + +// FillOption configures a Fill call. +type FillOption[K comparable] func(*fillOptions[K]) + +// WithFillScanBatch sets the number of pending documents fetched per scan. +// Values less than or equal to zero use 128. +func WithFillScanBatch[K comparable](size int) FillOption[K] { + return func(o *fillOptions[K]) { o.scanBatch = size } +} + +// WithFillSplit controls how each document's content is windowed into chunks. +func WithFillSplit[K comparable](options SplitOptions) FillOption[K] { + return func(o *fillOptions[K]) { o.split = options } +} + +// WithFillBatch controls how chunks are grouped into encode calls. A positive +// WithBatchSize packs chunks across documents within each scan page. Omitting +// it preserves the per-document encode unit. +func WithFillBatch[K comparable](options ...BatchOption) FillOption[K] { + return func(o *fillOptions[K]) { o.batch = applyBatchOptions(options) } +} + +// WithFillConcurrency bounds parallel fill work within each scan page. Values +// less than or equal to zero use one worker. SaveVectors and the error hooks +// stay serialized on the calling goroutine. When WithBatchSize is positive, +// this limit composes with WithBatchConcurrency; their product bounds the +// concurrent encode calls. +func WithFillConcurrency[K comparable](concurrency int) FillOption[K] { + return func(o *fillOptions[K]) { o.concurrency = concurrency } +} + +// WithFillEncodeError handles a document that fails to encode. Returning true +// skips and stamp-saves that document so later Fill calls do not retry it. +// Returning false aborts Fill, which is the right default for transient +// failures. Omitting this option also aborts on the first encode error. +func WithFillEncodeError[K comparable](handler func(doc K, err error) bool) FillOption[K] { + return func(o *fillOptions[K]) { o.onEncodeError = handler } +} + +// WithFillBatchErrorIsolation decides whether a failed shared encode call is +// worth diagnosing at document-slice granularity. Returning true permits +// diagnosis only; WithFillEncodeError still decides whether to skip an +// attributed document. Fill bypasses this handler for single-document calls, +// context errors, and errors with an exact *InvalidVectorError position. +func WithFillBatchErrorIsolation[K comparable](classify func(error) bool) FillOption[K] { + return func(o *fillOptions[K]) { o.shouldIsolateBatchError = classify } +} + +func applyFillOptions[K comparable](options []FillOption[K]) fillOptions[K] { + o := fillOptions[K]{} + for _, option := range options { + if option != nil { + option(&o) + } + } + return o } // FillStats reports what a Fill run embedded. @@ -54,7 +80,7 @@ type FillStats struct { // Chunks is the total chunk vectors saved across Documents. Chunks int // Skipped counts documents stamped without vectors because - // OnEncodeError elected to skip them. + // the WithFillEncodeError handler elected to skip them. Skipped int // Stale counts documents left pending because they changed between // scan and save (SaveVectors returned ErrStale). A later run retries @@ -72,14 +98,25 @@ type FillStats struct { // pending and not retried until the next Fill call, so an actively edited // document cannot starve the loop. // -// When Batch.BatchSize is positive, chunks from adjacent documents in one -// scan page may share an encode call. Errors with exact document attribution -// go directly to OnEncodeError. Other shared-call errors are diagnosed at -// document-slice granularity only when ShouldIsolateBatchError permits it; -// the nil default aborts without document-level retries. OnEncodeError remains -// the sole authority for skip-stamping an attributed document. -func Fill[K, G comparable](ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, o FillOptions[K]) (FillStats, error) { - scanBatch := o.ScanBatch +// When WithFillBatch includes a positive WithBatchSize, chunks from adjacent +// documents in one scan page may share an encode call. Errors with exact +// document attribution go directly to the WithFillEncodeError handler. Other +// shared-call errors are diagnosed at document-slice granularity only when +// WithFillBatchErrorIsolation permits it. Omitting that option aborts without +// document-level retries. WithFillEncodeError remains the sole authority for +// skip-stamping an attributed document. +func Fill[K, G comparable]( + ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, + options ...FillOption[K], +) (FillStats, error) { + o := applyFillOptions(options) + if err := ctx.Err(); err != nil { + return FillStats{}, err + } + if err := o.batch.validate(); err != nil { + return FillStats{}, err + } + scanBatch := o.scanBatch if scanBatch <= 0 { scanBatch = 128 } @@ -143,15 +180,15 @@ type fillBatchResult struct { err error } -// fillPage embeds one scan page. A positive BatchSize packs fixed-size chunk +// fillPage embeds one scan page. A positive batch size packs fixed-size chunk // batches across documents, then scatters vectors back before saving each // document independently. The legacy per-document path remains in use when -// BatchSize is unset, preserving its unbounded-per-document call shape. +// the batch size is unset, preserving its unbounded-per-document call shape. func fillPage[K, G comparable]( ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, - o FillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, + o fillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, ) error { - if o.Batch.BatchSize <= 0 || enc == nil { + if o.batch.batchSize <= 0 || enc == nil { return fillPageByDocument(ctx, store, gen, enc, o, docs, stale, stats) } return fillPageAcrossDocuments(ctx, store, gen, enc, o, docs, stale, stats) @@ -159,12 +196,12 @@ func fillPage[K, G comparable]( func fillPageAcrossDocuments[K, G comparable]( ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, - o FillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, + o fillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, ) error { documentStates := make([]fillDocumentState[K], len(docs)) var refs []fillChunkRef for doc, pending := range docs { - chunks := Split(pending.Content, o.Split) + chunks := Split(pending.Content, o.split) documentStates[doc] = fillDocumentState[K]{ encoded: fillEncoded[K]{ doc: pending, @@ -182,13 +219,14 @@ func fillPageAcrossDocuments[K, G comparable]( return saveReadyDocuments(ctx, store, gen, o, documentStates, true, stale, stats) } - orderedSaves := o.Concurrency <= 1 + batchSize := o.batch.effectiveBatchSize(len(refs)) + orderedSaves := o.concurrency <= 1 if err := saveReadyDocuments(ctx, store, gen, o, documentStates, orderedSaves, stale, stats); err != nil { return err } - batchConcurrency := max(o.Batch.Concurrency, 1) - batches := splitFillRefs(refs, o.Batch.BatchSize) + batchConcurrency := max(o.batch.concurrency, 1) + batches := splitFillRefs(refs, batchSize) encode := func(workCtx context.Context, batch []fillChunkRef) fillBatchResult { return encodeFillBatch(workCtx, enc, batch) } @@ -222,10 +260,10 @@ func fillPageAcrossDocuments[K, G comparable]( // its documents; applyFillBatch discards those results without racing // workers against the serialized document state. workers := len(batches) - // Compute min(Concurrency*Batch.Concurrency, len(batches)) without + // Compute the worker and batch concurrency product without // overflowing the product. - if o.Concurrency <= len(batches)/batchConcurrency { - workers = o.Concurrency * batchConcurrency + if o.concurrency <= len(batches)/batchConcurrency { + workers = o.concurrency * batchConcurrency } err = runFillJobs(ctx, workers, batches, encode, func(result fillBatchResult) bool { return result.err != nil }, @@ -266,7 +304,7 @@ func encodeFillBatch(ctx context.Context, enc EncodeFunc, refs []fillChunkRef) f for i, ref := range refs { chunks[i] = ref.value } - vectors, err := EncodeBatched(ctx, enc, chunks, BatchOptions{}) + vectors, err := EncodeBatched(ctx, enc, chunks) return fillBatchResult{refs: refs, vectors: vectors, err: err} } @@ -314,7 +352,7 @@ func splitFillRefsByDocument(refs []fillChunkRef) [][]fillChunkRef { } func decideFillDocumentError[K comparable]( - ctx context.Context, o FillOptions[K], states []fillDocumentState[K], doc int, err error, + ctx context.Context, o fillOptions[K], states []fillDocumentState[K], doc int, err error, ) error { state := &states[doc] if state.saved || state.failed { @@ -326,7 +364,7 @@ func decideFillDocumentError[K comparable]( if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return fmt.Errorf("encode document %v: %w", state.encoded.doc.Doc, err) } - if o.OnEncodeError == nil || !o.OnEncodeError(state.encoded.doc.Doc, err) { + if o.onEncodeError == nil || !o.onEncodeError(state.encoded.doc.Doc, err) { return fmt.Errorf("encode document %v: %w", state.encoded.doc.Doc, err) } state.encoded.err = err @@ -337,7 +375,7 @@ func decideFillDocumentError[K comparable]( } func probeFillDocumentSlices[K, G comparable]( - ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, o FillOptions[K], + ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, o fillOptions[K], refs []fillChunkRef, states []fillDocumentState[K], orderedSaves bool, stale map[K]struct{}, stats *FillStats, ) (bool, error) { @@ -375,7 +413,7 @@ func probeFillDocumentSlices[K, G comparable]( } func applyAttributedInvalidFillBatch[K, G comparable]( - ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, o FillOptions[K], + ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, o fillOptions[K], batch fillBatchResult, invalid *InvalidVectorError, states []fillDocumentState[K], orderedSaves bool, stale map[K]struct{}, stats *FillStats, ) error { @@ -409,7 +447,7 @@ func applyAttributedInvalidFillBatch[K, G comparable]( } func applyFillBatch[K, G comparable]( - ctx context.Context, store Store[K, G], gen G, o FillOptions[K], enc EncodeFunc, + ctx context.Context, store Store[K, G], gen G, o fillOptions[K], enc EncodeFunc, batch fillBatchResult, states []fillDocumentState[K], orderedSaves bool, stale map[K]struct{}, stats *FillStats, ) error { @@ -445,7 +483,7 @@ func applyFillBatch[K, G comparable]( if len(active) == 0 { return nil } - if o.ShouldIsolateBatchError == nil || !o.ShouldIsolateBatchError(batch.err) { + if o.shouldIsolateBatchError == nil || !o.shouldIsolateBatchError(batch.err) { return fillBatchContextError(batch.err, active, states) } failed, err := probeFillDocumentSlices(ctx, store, gen, enc, o, active, states, @@ -474,7 +512,7 @@ func fillBatchContextError[K comparable]( } func saveReadyDocuments[K, G comparable]( - ctx context.Context, store Store[K, G], gen G, o FillOptions[K], + ctx context.Context, store Store[K, G], gen G, o fillOptions[K], states []fillDocumentState[K], ordered bool, stale map[K]struct{}, stats *FillStats, ) error { for i := range states { @@ -594,18 +632,18 @@ func runFillJobs[J, R any]( return ctx.Err() } -// fillPageByDocument is the original per-document path used when BatchSize is -// unset. Workers split and encode up to o.Concurrency documents in parallel +// fillPageByDocument is the original per-document path used when the batch +// size is unset. Workers split and encode up to o.concurrency documents in parallel // while the calling goroutine saves each result as it completes. The first // save-side failure cancels in-flight encodes and is returned. func fillPageByDocument[K, G comparable]( ctx context.Context, store Store[K, G], gen G, enc EncodeFunc, - o FillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, + o fillOptions[K], docs []Pending[K], stale map[K]struct{}, stats *FillStats, ) error { - return runFillJobs(ctx, o.Concurrency, docs, + return runFillJobs(ctx, o.concurrency, docs, func(workCtx context.Context, p Pending[K]) fillEncoded[K] { - chunks := Split(p.Content, o.Split) - vectors, err := EncodeBatched(workCtx, enc, chunks, o.Batch) + chunks := Split(p.Content, o.split) + vectors, err := encodeBatched(workCtx, enc, chunks, o.batch) return fillEncoded[K]{doc: p, chunks: chunks, vectors: vectors, err: err} }, nil, @@ -614,12 +652,12 @@ func fillPageByDocument[K, G comparable]( }) } -// saveEncoded applies one document's encode outcome: it consults -// OnEncodeError for failures, stamps skips, saves vectors, and records +// saveEncoded applies one document's encode outcome: it consults the encode +// error handler for failures, stamps skips, saves vectors, and records // stale revisions, updating stats to match. It runs on Fill's calling // goroutine only. func saveEncoded[K, G comparable]( - ctx context.Context, store Store[K, G], gen G, o FillOptions[K], + ctx context.Context, store Store[K, G], gen G, o fillOptions[K], r fillEncoded[K], stale map[K]struct{}, stats *FillStats, ) error { skipped := r.skip @@ -630,7 +668,7 @@ func saveEncoded[K, G comparable]( if errors.Is(r.err, context.Canceled) || errors.Is(r.err, context.DeadlineExceeded) { return fmt.Errorf("encode document %v: %w", r.doc.Doc, r.err) } - if o.OnEncodeError == nil || !o.OnEncodeError(r.doc.Doc, r.err) { + if o.onEncodeError == nil || !o.onEncodeError(r.doc.Doc, r.err) { return fmt.Errorf("encode document %v: %w", r.doc.Doc, r.err) } skipped = true @@ -696,7 +734,7 @@ func Search[K, G comparable]( if enc == nil { return nil, fmt.Errorf("no encoder for generation %v", gen) } - vectors, err := EncodeBatched(ctx, enc, []Chunk{{Index: 0, Text: queryText}}, BatchOptions{}) + vectors, err := EncodeBatched(ctx, enc, []Chunk{{Index: 0, Text: queryText}}) if err != nil { return nil, fmt.Errorf("embed query for generation %v: %w", gen, err) } diff --git a/vector/flow_batch_error_test.go b/vector/flow_batch_error_test.go index 513cc1f..cb0b078 100644 --- a/vector/flow_batch_error_test.go +++ b/vector/flow_batch_error_test.go @@ -54,15 +54,15 @@ func TestFillSharedErrorClassifierFailsClosedWithoutProbes(t *testing.T) { return tc.classifier(err) } } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3}, - ShouldIsolateBatchError: classifier, - OnEncodeError: func(int64, error) bool { + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64](vector.WithBatchSize(3)), + vector.WithFillBatchErrorIsolation[int64](classifier), + vector.WithFillEncodeError[int64](func(int64, error) bool { hooks++ return true - }, - }) + }), + ) require.Error(err) var got *fillProviderError require.ErrorAs(err, &got) @@ -99,15 +99,15 @@ func TestFillSharedErrorRejectedFirstProbeStopsDiagnosis(t *testing.T) { return tc.hook(doc, err) } } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++ return true - }, - OnEncodeError: hook, - }) + }), + vector.WithFillEncodeError[int64](hook), + ) require.Error(err) var providerErr *fillProviderError require.ErrorAs(err, &providerErr) @@ -128,15 +128,15 @@ func TestFillSharedErrorAllowsTwoPoisonDocuments(t *testing.T) { calls++ return nil, &fillProviderError{code: 400} } - stats, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(doc int64, _ error) bool { + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { hooks[doc]++ return true - }, - }) + }), + ) Require.NoError(t, err) assert.Equal(3, calls) assert.Equal(map[int64]int{1: 1, 2: 1}, hooks) @@ -153,22 +153,22 @@ func TestFillSharedInvalidVectorRejectedWithoutProbe(t *testing.T) { calls++ return [][]float32{{1}, {0}, {1}}, nil } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3}, - ShouldIsolateBatchError: func(error) bool { + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64](vector.WithBatchSize(3)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++ return true - }, - OnEncodeError: func(doc int64, err error) bool { + }), + vector.WithFillEncodeError[int64](func(doc int64, err error) bool { hooks++ assert.Equal(int64(2), doc) var invalid *vector.InvalidVectorError Require.ErrorAs(t, err, &invalid) assert.Equal(0, invalid.Chunk) return false - }, - }) + }), + ) Require.Error(t, err) assert.Equal(1, calls) assert.Zero(classifiers) @@ -183,11 +183,11 @@ func TestFillSharedInvalidVectorNilHookRejectsWithoutProbe(t *testing.T) { calls++ return [][]float32{{1}, {0}}, nil } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { classifiers++; return true }, - }) + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++; return true }), + ) Require.Error(t, err) Assert.Equal(t, 1, calls) Assert.Zero(t, classifiers) @@ -210,14 +210,14 @@ func TestFillSharedInvalidVectorRecoversOnlyOtherSlices(t *testing.T) { } return out, nil } - stats, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3}, - OnEncodeError: func(doc int64, _ error) bool { + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64](vector.WithBatchSize(3)), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { hooks[doc]++ return true - }, - }) + }), + ) Require.NoError(t, err) assert.Equal([][]string{{"good", "bad", "later"}, {"good"}, {"later"}}, calls) assert.Equal(map[int64]int{2: 1}, hooks) @@ -239,15 +239,15 @@ func TestFillSharedInvalidRecoveryFailureUsesProbeRules(t *testing.T) { } return nil, &fillProviderError{code: 400} } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { classifiers++; return true }, - OnEncodeError: func(doc int64, _ error) bool { + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++; return true }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { hooks[doc]++ return doc == 1 - }, - }) + }), + ) require.Error(err) var providerErr *fillProviderError require.ErrorAs(err, &providerErr) @@ -266,12 +266,12 @@ func TestFillSharedInvalidVectorOutOfRangeIsFatal(t *testing.T) { calls++ return nil, &vector.InvalidVectorError{Chunk: 2, Component: -1, Reason: "zero norm"} } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { classifiers++; return true }, - OnEncodeError: func(int64, error) bool { hooks++; return true }, - }) + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++; return true }), + vector.WithFillEncodeError[int64](func(int64, error) bool { hooks++; return true }), + ) require.ErrorContains(err, "invalid vector chunk 2 outside batch of 2 chunks") var invalid *vector.InvalidVectorError require.ErrorAs(err, &invalid) @@ -294,10 +294,10 @@ func TestFillSharedInvalidVectorPreservesCompanionCauses(t *testing.T) { sentinel, ) } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - OnEncodeError: func(doc int64, err error) bool { + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillEncodeError[int64](func(doc int64, err error) bool { assert.Equal(int64(2), doc) var invalid *vector.InvalidVectorError require.ErrorAs(err, &invalid) @@ -307,8 +307,8 @@ func TestFillSharedInvalidVectorPreservesCompanionCauses(t *testing.T) { assert.Same(providerErr, gotProvider) assert.ErrorIs(err, sentinel) return false - }, - }) + }), + ) require.Error(err) var invalid *vector.InvalidVectorError require.ErrorAs(err, &invalid) @@ -367,13 +367,13 @@ func TestFillRejectedProbeBackpressuresAndCancelsWorkers(t *testing.T) { fillReturned := make(chan struct{}) go func() { defer close(fillReturned) - _, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 6, - Batch: vector.BatchOptions{BatchSize: 2}, - Concurrency: 2, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(int64, error) bool { return false }, - }) + _, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](6), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillConcurrency[int64](2), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(int64, error) bool { return false }), + ) done <- err }() t.Cleanup(func() { @@ -434,22 +434,22 @@ func TestFillLateSharedFailureFiltersDecidedDocument(t *testing.T) { return nil, fmt.Errorf("unexpected texts %q", texts) } } - stats, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Split: vector.SplitOptions{MaxRunes: 1}, - Batch: vector.BatchOptions{BatchSize: 2}, - Concurrency: 2, - ShouldIsolateBatchError: func(error) bool { + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 1}), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillConcurrency[int64](2), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifierCalls.Add(1) return true - }, - OnEncodeError: func(doc int64, _ error) bool { + }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { hookCalls.Add(1) assert.Equal(int64(1), doc) close(releaseShared) return true - }, - }) + }), + ) Require.NoError(t, err) assert.Equal(int32(1), hookCalls.Load()) assert.Equal(int32(1), classifierCalls.Load()) @@ -469,12 +469,12 @@ func TestFillWrappedProbeDeadlineAbortsWithoutHook(t *testing.T) { } return nil, fmt.Errorf("encoder timeout: %w", context.DeadlineExceeded) } - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(int64, error) bool { hooks++; return true }, - }) + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(int64, error) bool { hooks++; return true }), + ) Require.ErrorIs(t, err, context.DeadlineExceeded) Assert.Equal(t, 2, calls) Assert.Zero(t, hooks) @@ -510,12 +510,11 @@ func TestFillBatchClassifierExclusions(t *testing.T) { var classifiers, hooks int _, err := vector.Fill(context.Background(), store, 7, func(context.Context, []string) ([][]float32, error) { return nil, tc.encodeErr }, - vector.FillOptions[int64]{ - ScanBatch: len(tc.content), - Batch: vector.BatchOptions{BatchSize: tc.batchSize}, - ShouldIsolateBatchError: func(error) bool { classifiers++; return true }, - OnEncodeError: func(int64, error) bool { hooks++; return true }, - }) + vector.WithFillScanBatch[int64](len(tc.content)), + vector.WithFillBatch[int64](vector.WithBatchSize(tc.batchSize)), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { classifiers++; return true }), + vector.WithFillEncodeError[int64](func(int64, error) bool { hooks++; return true }), + ) if errors.Is(tc.encodeErr, context.Canceled) { require.ErrorIs(err, context.Canceled) } else { diff --git a/vector/flow_internal_test.go b/vector/flow_internal_test.go index e9427dd..d98279c 100644 --- a/vector/flow_internal_test.go +++ b/vector/flow_internal_test.go @@ -70,13 +70,13 @@ func TestApplyFillBatchProbeInvalidVectorAddsSliceAndLocalOffsets(t *testing.T) batch := encodeFillBatch(context.Background(), enc, refs) var got *InvalidVectorError err := applyFillBatch(context.Background(), noOpFillStore{}, 7, - FillOptions[int64]{ - ShouldIsolateBatchError: func(err error) bool { + fillOptions[int64]{ + shouldIsolateBatchError: func(err error) bool { var providerErr *internalFillProviderError require.ErrorAs(err, &providerErr) return true }, - OnEncodeError: func(doc int64, err error) bool { + onEncodeError: func(doc int64, err error) bool { assert.Equal(int64(10), doc) require.ErrorAs(err, &got) return false @@ -133,13 +133,13 @@ func TestApplyFillBatchProbeInvalidVectorPreservesCompanionCauses(t *testing.T) batch := encodeFillBatch(context.Background(), enc, refs) var gotInvalid *InvalidVectorError err := applyFillBatch(context.Background(), noOpFillStore{}, 7, - FillOptions[int64]{ - ShouldIsolateBatchError: func(err error) bool { + fillOptions[int64]{ + shouldIsolateBatchError: func(err error) bool { var providerErr *internalFillProviderError require.ErrorAs(err, &providerErr) return true }, - OnEncodeError: func(doc int64, err error) bool { + onEncodeError: func(doc int64, err error) bool { assert.Equal(int64(10), doc) require.ErrorAs(err, &gotInvalid) assert.Equal(4, gotInvalid.Chunk) diff --git a/vector/flow_test.go b/vector/flow_test.go index a3140c6..411279e 100644 --- a/vector/flow_test.go +++ b/vector/flow_test.go @@ -139,10 +139,10 @@ func TestFillEmbedsAllPendingThenStops(t *testing.T) { store.content[1] = "alpha" store.content[2] = "beta gamma delta" - stats, err := vector.Fill(ctx, store, 7, lenEncoder(), vector.FillOptions[int64]{ - ScanBatch: 1, // force multiple scan rounds - Split: vector.SplitOptions{MaxRunes: 4, Overlap: 0}, - }) + stats, err := vector.Fill(ctx, store, 7, lenEncoder(), + vector.WithFillScanBatch[int64](1), // force multiple scan rounds + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 4, Overlap: 0}), + ) require.NoError(err) assert.Equal(2, stats.Documents) @@ -151,12 +151,12 @@ func TestFillEmbedsAllPendingThenStops(t *testing.T) { assert.InDelta(4, store.vectors[7][1][0].Vector[0], 1e-6, "first chunk carries its rune length") // A second run finds nothing pending and embeds zero documents. - again, err := vector.Fill(ctx, store, 7, lenEncoder(), vector.FillOptions[int64]{}) + again, err := vector.Fill(ctx, store, 7, lenEncoder()) require.NoError(err) assert.Equal(0, again.Documents) } -func TestFillBatchesChunksAcrossDocuments(t *testing.T) { +func TestFillBatchesChunksAcrossDocumentsWithinTokenBudget(t *testing.T) { require := require.New(t) assert := assert.New(t) ctx := context.Background() @@ -172,15 +172,19 @@ func TestFillBatchesChunksAcrossDocuments(t *testing.T) { return lenEncoder()(ctx, texts) } - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 7, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - }) + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](7), + vector.WithFillBatch[int64]( + vector.WithBatchSize(7), + vector.WithBatchConcurrency(1), + vector.WithBatchTokenBudget(3, 1), + ), + vector.WithFillConcurrency[int64](1), + ) require.NoError(err) assert.Equal([]int{3, 3, 1}, batchSizes, - "BatchSize should pack chunks from adjacent documents into each encode call") + "the token budget caps batches below the count limit") assert.Equal(7, stats.Documents) assert.Equal(7, stats.Chunks) for doc := int64(1); doc <= 7; doc++ { @@ -190,6 +194,61 @@ func TestFillBatchesChunksAcrossDocuments(t *testing.T) { } } +func TestFillRejectsInputAboveTokenBudgetBeforeEncoder(t *testing.T) { + assert := assert.New(t) + require := require.New(t) + store := newMemStore() + store.content[1] = "one" + var calls atomic.Int64 + enc := func(context.Context, []string) ([][]float32, error) { + calls.Add(1) + return [][]float32{{1}}, nil + } + + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillBatch[int64]( + vector.WithBatchSize(1), + vector.WithBatchTokenBudget(31_999, 32_000), + ), + ) + + require.Error(err) + assert.ErrorContains(err, "token budget") + assert.Zero(calls.Load(), "an invalid budget is rejected before the provider call") + assert.Zero(stats.Documents) + assert.False(store.embedded[1][7]) +} + +func TestFillDoesNotSkipInvalidTokenBudgetWithoutBatchSize(t *testing.T) { + assert := assert.New(t) + require := require.New(t) + store := newMemStore() + store.content[1] = "one" + var calls, hookCalls atomic.Int64 + enc := func(context.Context, []string) ([][]float32, error) { + calls.Add(1) + return [][]float32{{1}}, nil + } + + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillBatch[int64]( + vector.WithBatchTokenBudget(31_999, 32_000), + ), + vector.WithFillEncodeError[int64](func(int64, error) bool { + hookCalls.Add(1) + return true + }), + ) + + require.Error(err) + assert.ErrorContains(err, "token budget") + assert.Zero(calls.Load(), "an invalid budget is rejected before the provider call") + assert.Zero(hookCalls.Load(), "configuration errors bypass the document error handler") + assert.Zero(stats.Documents) + assert.Zero(stats.Skipped) + assert.False(store.embedded[1][7]) +} + func TestFillCrossDocumentBatchingMatchesPerDocumentVectors(t *testing.T) { require := require.New(t) assert := assert.New(t) @@ -206,11 +265,10 @@ func TestFillCrossDocumentBatchingMatchesPerDocumentVectors(t *testing.T) { baseline.content[doc] = content batched.content[doc] = content } - options := vector.FillOptions[int64]{ - ScanBatch: 3, - Split: vector.SplitOptions{MaxRunes: 2}, - } - baselineStats, err := vector.Fill(ctx, baseline, 7, textEncoder(), options) + baselineStats, err := vector.Fill(ctx, baseline, 7, textEncoder(), + vector.WithFillScanBatch[int64](3), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 2}), + ) require.NoError(err) var batchSizes []int @@ -218,9 +276,13 @@ func TestFillCrossDocumentBatchingMatchesPerDocumentVectors(t *testing.T) { batchSizes = append(batchSizes, len(texts)) return textEncoder()(ctx, texts) } - options.Batch = vector.BatchOptions{BatchSize: 3, Concurrency: 1} - options.Concurrency = 1 - batchedStats, err := vector.Fill(ctx, batched, 7, batchedEncoder, options) + batchedStats, err := vector.Fill(ctx, batched, 7, batchedEncoder, + vector.WithFillScanBatch[int64](3), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 2}), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + ) require.NoError(err) assert.Equal([]int{3, 3, 2}, batchSizes, @@ -242,10 +304,10 @@ func TestFillCrossDocumentBatchingMatchesLegacyAcrossConfigurations(t *testing.T baseline := newMemStore() maps.Copy(baseline.content, contents) - baselineStats, err := vector.Fill(ctx, baseline, 7, textEncoder(), vector.FillOptions[int64]{ - ScanBatch: 5, - Split: vector.SplitOptions{MaxRunes: 3, Overlap: 1}, - }) + baselineStats, err := vector.Fill(ctx, baseline, 7, textEncoder(), + vector.WithFillScanBatch[int64](5), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 3, Overlap: 1}), + ) require.NoError(t, err) for _, scanBatch := range []int{1, 2, 5} { @@ -258,15 +320,15 @@ func TestFillCrossDocumentBatchingMatchesLegacyAcrossConfigurations(t *testing.T store := newMemStore() maps.Copy(store.content, contents) - stats, err := vector.Fill(ctx, store, 7, textEncoder(), vector.FillOptions[int64]{ - ScanBatch: scanBatch, - Split: vector.SplitOptions{MaxRunes: 3, Overlap: 1}, - Batch: vector.BatchOptions{ - BatchSize: batchSize, - Concurrency: batchConcurrency, - }, - Concurrency: fillConcurrency, - }) + stats, err := vector.Fill(ctx, store, 7, textEncoder(), + vector.WithFillScanBatch[int64](scanBatch), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 3, Overlap: 1}), + vector.WithFillBatch[int64]( + vector.WithBatchSize(batchSize), + vector.WithBatchConcurrency(batchConcurrency), + ), + vector.WithFillConcurrency[int64](fillConcurrency), + ) require.NoError(t, err) assert.Equal(t, baselineStats, stats) assert.Equal(t, baseline.vectors, store.vectors) @@ -294,16 +356,17 @@ func TestFillCrossDocumentBatchingIsolatesPoisonDocument(t *testing.T) { return base(ctx, texts) } var skipped []int64 - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(doc int64, _ error) bool { + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { skipped = append(skipped, doc) return true - }, - }) + }), + ) require.NoError(err) assert.Equal([]int{3, 1, 1, 1}, batchSizes, @@ -337,17 +400,18 @@ func TestFillCrossDocumentBatchingTranslatesInvalidVectorChunkIndex(t *testing.T } var gotInvalid *vector.InvalidVectorError - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Split: vector.SplitOptions{MaxRunes: 1}, - Batch: vector.BatchOptions{BatchSize: 2, Concurrency: 1}, - Concurrency: 1, - OnEncodeError: func(doc int64, err error) bool { + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillSplit[int64](vector.SplitOptions{MaxRunes: 1}), + vector.WithFillBatch[int64]( + vector.WithBatchSize(2), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + vector.WithFillEncodeError[int64](func(doc int64, err error) bool { assert.Equal(int64(2), doc) require.ErrorAs(err, &gotInvalid) return true - }, - }) + }), + ) require.NoError(err) require.NotNil(gotInvalid) @@ -377,11 +441,12 @@ func TestFillCrossDocumentBatchingLeavesOnlyChangedDocumentPending(t *testing.T) return out, nil } - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2, Concurrency: 1}, - Concurrency: 1, - }) + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64]( + vector.WithBatchSize(2), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + ) require.NoError(err) assert.Equal(1, stats.Documents) @@ -405,11 +470,12 @@ func TestFillCrossDocumentBatchingStampsBlankDocuments(t *testing.T) { return lenEncoder()(ctx, texts) } - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - }) + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + ) require.NoError(err) assert.Equal(int64(1), calls.Load()) @@ -429,10 +495,10 @@ func TestFillCrossDocumentBatchingRejectsNilEncoderBeforeStampingEmptyDocuments( store.content[1] = "" store.content[2] = "" - stats, err := vector.Fill(context.Background(), store, 7, nil, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2}, - }) + stats, err := vector.Fill(context.Background(), store, 7, nil, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64](vector.WithBatchSize(2)), + ) require.Error(err) assert.Zero(stats.Documents) @@ -450,16 +516,17 @@ func TestFillCrossDocumentBatchingEncodeErrorAbortsAtFailedDocument(t *testing.T store.content[2] = "poison" store.content[3] = "fine two" var consulted []int64 - _, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(doc int64, _ error) bool { + _, err := vector.Fill(ctx, store, 7, poisonEncoder(), + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { consulted = append(consulted, doc) return false - }, - }) + }), + ) require.ErrorContains(err, "encode document 2") assert.Equal([]int64{2}, consulted) @@ -486,16 +553,17 @@ func TestFillCrossDocumentBatchingAbortsUnattributedBatchError(t *testing.T) { return [][]float32{{1}}, nil } called := false - _, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 3, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(int64, error) bool { + _, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](3), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(int64, error) bool { called = true return true - }, - }) + }), + ) require.ErrorContains(err, "no document failed in isolation") assert.Equal([]int{3, 1, 1, 1}, batchSizes) @@ -518,15 +586,16 @@ func TestFillCrossDocumentBatchingDoesNotSkipCancelledEncode(t *testing.T) { enc := func(context.Context, []string) ([][]float32, error) { return nil, context.Canceled } - _, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 2, Concurrency: 1}, - Concurrency: 1, - OnEncodeError: func(int64, error) bool { + _, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64]( + vector.WithBatchSize(2), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + vector.WithFillEncodeError[int64](func(int64, error) bool { called = true return true - }, - }) + }), + ) require.ErrorIs(err, context.Canceled) assert.False(called, "cancellation bypasses the permanent-error hook") @@ -557,7 +626,7 @@ func TestFillLeavesChangedDocumentPending(t *testing.T) { return out, nil } - stats, err := vector.Fill(ctx, store, 7, racingEnc, vector.FillOptions[int64]{}) + stats, err := vector.Fill(ctx, store, 7, racingEnc) require.NoError(err) assert.Equal(1, stats.Documents, "the unchanged doc is embedded") assert.Equal(1, stats.Stale, "the changed doc is reported stale") @@ -565,7 +634,7 @@ func TestFillLeavesChangedDocumentPending(t *testing.T) { assert.True(store.embedded[2][7]) // The next run re-reads the document at its new revision and succeeds. - again, err := vector.Fill(ctx, store, 7, lenEncoder(), vector.FillOptions[int64]{}) + again, err := vector.Fill(ctx, store, 7, lenEncoder()) require.NoError(err) assert.Equal(1, again.Documents) assert.True(store.embedded[1][7]) @@ -581,12 +650,12 @@ func TestFillSkipHookStampsFailedDocument(t *testing.T) { store.content[2] = "fine" var skipped []int64 - stats, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{ - OnEncodeError: func(doc int64, err error) bool { + stats, err := vector.Fill(ctx, store, 7, poisonEncoder(), + vector.WithFillEncodeError[int64](func(doc int64, err error) bool { skipped = append(skipped, doc) return true - }, - }) + }), + ) require.NoError(err) assert.Equal(1, stats.Documents) assert.Equal(1, stats.Skipped) @@ -595,7 +664,7 @@ func TestFillSkipHookStampsFailedDocument(t *testing.T) { assert.Empty(store.vectors[7][1], "skipped doc has no vectors") assert.True(store.embedded[2][7]) - again, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{}) + again, err := vector.Fill(ctx, store, 7, poisonEncoder()) require.NoError(err) assert.Equal(0, again.Documents, "a stamped skip does not reappear as pending") } @@ -608,12 +677,11 @@ func TestFillEncodeErrorAbortsWithoutSkip(t *testing.T) { store := newMemStore() store.content[1] = "poison" - _, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{}) + _, err := vector.Fill(ctx, store, 7, poisonEncoder()) require.ErrorContains(err, "encode document") - _, err = vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{ - OnEncodeError: func(int64, error) bool { return false }, - }) + _, err = vector.Fill(ctx, store, 7, poisonEncoder(), + vector.WithFillEncodeError[int64](func(int64, error) bool { return false })) require.ErrorContains(err, "encode document") assert.False(store.embedded[1][7], "an aborted doc is neither embedded nor stamped") } @@ -630,12 +698,12 @@ func TestFillDoesNotSkipCancelledEncode(t *testing.T) { enc := func(context.Context, []string) ([][]float32, error) { return nil, context.Canceled } - _, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - OnEncodeError: func(int64, error) bool { + _, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillEncodeError[int64](func(int64, error) bool { called = true return true - }, - }) + }), + ) require.ErrorIs(err, context.Canceled) assert.False(called, "cancellation bypasses the permanent-error skip hook") assert.False(store.embedded[1][7], "a cancelled document is not stamped as handled") @@ -673,9 +741,8 @@ func TestFillConcurrencyEncodesDocumentsInParallel(t *testing.T) { return out, nil } - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - Concurrency: workers, - }) + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillConcurrency[int64](workers)) require.NoError(err) assert.Equal(workers, stats.Documents) @@ -761,7 +828,7 @@ func TestFillDefaultConcurrencyIsSequential(t *testing.T) { inSave.Store(false) }} - stats, err := vector.Fill(ctx, hooked, 7, enc, vector.FillOptions[int64]{}) + stats, err := vector.Fill(ctx, hooked, 7, enc) require.NoError(err) require.Equal(6, stats.Documents) require.False(overlapped.Load(), @@ -785,11 +852,12 @@ func TestFillCrossDocumentBatchingDoesNotEncodeNextWindowAfterSaveFailure(t *tes return lenEncoder()(ctx, texts) } - _, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 7, - Batch: vector.BatchOptions{BatchSize: 3, Concurrency: 1}, - Concurrency: 1, - }) + _, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](7), + vector.WithFillBatch[int64]( + vector.WithBatchSize(3), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](1), + ) require.ErrorIs(err, sentinel) assert.Equal(int64(1), calls.Load(), @@ -833,14 +901,12 @@ func TestFillCrossDocumentBatchingComposesConcurrencyBounds(t *testing.T) { return out, nil } - stats, err := vector.Fill(ctx, store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 8, - Batch: vector.BatchOptions{ - BatchSize: 2, - Concurrency: 2, - }, - Concurrency: 2, - }) + stats, err := vector.Fill(ctx, store, 7, enc, + vector.WithFillScanBatch[int64](8), + vector.WithFillBatch[int64]( + vector.WithBatchSize(2), vector.WithBatchConcurrency(2)), + vector.WithFillConcurrency[int64](2), + ) require.NoError(err) assert.Equal(int64(maxCalls), observedMax.Load(), @@ -889,15 +955,16 @@ func TestFillCrossDocumentBatchingConcurrentFailureKeepsAttribution(t *testing.T } done := make(chan fillResult, 1) go func() { - stats, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 4, - Batch: vector.BatchOptions{BatchSize: 2, Concurrency: 1}, - Concurrency: 2, - ShouldIsolateBatchError: func(error) bool { return true }, - OnEncodeError: func(doc int64, _ error) bool { + stats, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](4), + vector.WithFillBatch[int64]( + vector.WithBatchSize(2), vector.WithBatchConcurrency(1)), + vector.WithFillConcurrency[int64](2), + vector.WithFillBatchErrorIsolation[int64](func(error) bool { return true }), + vector.WithFillEncodeError[int64](func(doc int64, _ error) bool { return doc == 1 - }, - }) + }), + ) done <- fillResult{stats: stats, err: err} }() @@ -937,11 +1004,12 @@ func TestFillCrossDocumentBatchingDoesNotBlockCompletedSaves(t *testing.T) { done := make(chan error, 1) go func() { - _, err := vector.Fill(context.Background(), store, 7, enc, vector.FillOptions[int64]{ - ScanBatch: 2, - Batch: vector.BatchOptions{BatchSize: 1, Concurrency: 2}, - Concurrency: 2, - }) + _, err := vector.Fill(context.Background(), store, 7, enc, + vector.WithFillScanBatch[int64](2), + vector.WithFillBatch[int64]( + vector.WithBatchSize(1), vector.WithBatchConcurrency(2)), + vector.WithFillConcurrency[int64](2), + ) done <- err }() <-slowStarted @@ -971,13 +1039,13 @@ func TestFillConcurrencySkipHookStampsFailedDocument(t *testing.T) { store.content[3] = "also fine" var skipped []int64 - stats, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{ - Concurrency: 3, - OnEncodeError: func(doc int64, err error) bool { + stats, err := vector.Fill(ctx, store, 7, poisonEncoder(), + vector.WithFillConcurrency[int64](3), + vector.WithFillEncodeError[int64](func(doc int64, err error) bool { skipped = append(skipped, doc) return true - }, - }) + }), + ) require.NoError(err) assert.Equal(2, stats.Documents) assert.Equal(1, stats.Skipped) @@ -999,17 +1067,15 @@ func TestFillConcurrencyEncodeErrorAborts(t *testing.T) { store.content[doc] = "fine" } - _, err := vector.Fill(ctx, store, 7, poisonEncoder(), vector.FillOptions[int64]{ - Concurrency: 4, - }) + _, err := vector.Fill(ctx, store, 7, poisonEncoder(), + vector.WithFillConcurrency[int64](4)) require.ErrorContains(err, "encode document") assert.False(store.embedded[1][7], "the failed doc is neither embedded nor stamped") // The failed document stays pending: a later run with a working encoder // picks up everything the aborted page left behind. - again, err := vector.Fill(ctx, store, 7, lenEncoder(), vector.FillOptions[int64]{ - Concurrency: 4, - }) + again, err := vector.Fill(ctx, store, 7, lenEncoder(), + vector.WithFillConcurrency[int64](4)) require.NoError(err) assert.True(store.embedded[1][7]) assert.Equal(0, again.Stale) diff --git a/vector/sqlitevec/sqlitevec_test.go b/vector/sqlitevec/sqlitevec_test.go index 11d97ff..ffc8ab9 100644 --- a/vector/sqlitevec/sqlitevec_test.go +++ b/vector/sqlitevec/sqlitevec_test.go @@ -143,7 +143,7 @@ func TestStoreFillThenSearch(t *testing.T) { require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - stats, err := vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + stats, err := vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) assert.Equal(2, stats.Documents) @@ -187,7 +187,7 @@ func TestStoreSearchUnionsLiveGenerations(t *testing.T) { require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "v1", Dimensions: 3}, sqlitevec.StateActive)) - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) // The building generation has covered only doc 1 so far. @@ -564,7 +564,7 @@ func TestStoreFillWithRevisionColumn(t *testing.T) { require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - stats, err := vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + stats, err := vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) assert.Equal(2, stats.Documents) @@ -634,7 +634,7 @@ func TestStoreQueryGenerationExcludesDeletedDocuments(t *testing.T) { _, err := db.ExecContext(ctx, `INSERT INTO messages (id, body) VALUES (1, 'a cat sat'), (2, 'a dog ran')`) require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) // The caller deletes a source row without telling the store. @@ -660,7 +660,7 @@ func TestStoreQueryGenerationExcludesEditedDocumentUntilReembedded(t *testing.T) _, err := db.ExecContext(ctx, `INSERT INTO messages (id, body, last_modified) VALUES (1, 'a cat sat', 1)`) require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) // The caller redacts the content, bumping the revision. @@ -671,7 +671,7 @@ func TestStoreQueryGenerationExcludesEditedDocumentUntilReembedded(t *testing.T) require.NoError(err) assert.Empty(hits, "the pre-edit vector never surfaces after the revision changes") - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) hits, err = store.QueryGeneration(ctx, 1, vector.Vector{0, 1, 0}, 10) @@ -689,7 +689,7 @@ func TestStoreQueryGenerationExcludesInvalidatedDocumentUntilReembedded(t *testi _, err := db.ExecContext(ctx, `INSERT INTO messages (id, body) VALUES (1, 'a cat sat')`) require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) // Without a revision column, the caller signals an edit by clearing @@ -701,7 +701,7 @@ func TestStoreQueryGenerationExcludesInvalidatedDocumentUntilReembedded(t *testi require.NoError(err) assert.Empty(hits, "the pre-invalidation vector never surfaces while the document is pending") - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) hits, err = store.QueryGeneration(ctx, 1, vector.Vector{0, 1, 0}, 10) diff --git a/vector/sqlitevec/store_plan_test.go b/vector/sqlitevec/store_plan_test.go index 4ec3502..e36d6c6 100644 --- a/vector/sqlitevec/store_plan_test.go +++ b/vector/sqlitevec/store_plan_test.go @@ -29,7 +29,7 @@ func TestQueryGenerationPlanScansKNNOnce(t *testing.T) { _, err := db.ExecContext(ctx, `INSERT INTO messages (id, body) VALUES (1, 'a cat sat'), (2, 'a dog ran')`) require.NoError(err) require.NoError(store.EnsureGeneration(ctx, 1, vector.Generation{Model: "m", Dimensions: 3}, sqlitevec.StateActive)) - _, err = vector.Fill(ctx, store, 1, topicEncoder(), vector.FillOptions[int64]{}) + _, err = vector.Fill(ctx, store, 1, topicEncoder()) require.NoError(err) sqlText, args, err := store.QueryGenerationSQLForTest(1, []float32{1, 0, 0}, 10)