|
| 1 | +package vector |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "slices" |
| 7 | + "strings" |
| 8 | + "testing" |
| 9 | + |
| 10 | + "github.com/stretchr/testify/assert" |
| 11 | + "github.com/stretchr/testify/require" |
| 12 | + kitvec "go.kenn.io/kit/vector" |
| 13 | +) |
| 14 | + |
| 15 | +// blankWindowContent is long enough that its middle 4000-rune window (the |
| 16 | +// test index's MaxRunes, stride 3400 after the 15% overlap) holds nothing but |
| 17 | +// spaces, so kitvec.Split drops that window and numbers the surviving chunks |
| 18 | +// 0 and 2. Everything that re-derives chunks from content has to cope with |
| 19 | +// that gap. |
| 20 | +func blankWindowContent() string { |
| 21 | + return strings.Repeat("a", 100) + strings.Repeat(" ", 7500) + strings.Repeat("b", 100) |
| 22 | +} |
| 23 | + |
| 24 | +// recordingEncoder returns a unit-vector encoder that appends every text it is |
| 25 | +// asked to embed to seen. |
| 26 | +func recordingEncoder(seen *[]string) kitvec.EncodeFunc { |
| 27 | + return func(_ context.Context, texts []string) ([][]float32, error) { |
| 28 | + *seen = append(*seen, texts...) |
| 29 | + out := make([][]float32, len(texts)) |
| 30 | + for i := range texts { |
| 31 | + out[i] = []float32{1, 0, 0} |
| 32 | + } |
| 33 | + return out, nil |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +// TestBuildStampsWhitespaceOnlyDocumentWithoutEmbeddingIt covers the whole |
| 38 | +// point of the kit 0.13 upgrade: a document whose content is only whitespace |
| 39 | +// is stamped for the generation with no vectors and never becomes an |
| 40 | +// embeddings request, so no provider ever gets the chance to reject it. The |
| 41 | +// build still completes and auto-activates. |
| 42 | +func TestBuildStampsWhitespaceOnlyDocumentWithoutEmbeddingIt(t *testing.T) { |
| 43 | + ix := openTestIndex(t) |
| 44 | + ctx := context.Background() |
| 45 | + src := &fakeUnitSource{rows: []fakeUnit{ |
| 46 | + {unit: userDoc("s1", "", 0, " \n\t "), endedAt: "2024-01-01T00:00:00Z"}, |
| 47 | + {unit: userDoc("s1", "", 1, "real content"), endedAt: "2024-01-01T00:00:01Z"}, |
| 48 | + }} |
| 49 | + |
| 50 | + var seen []string |
| 51 | + result, err := ix.Build(ctx, src, recordingEncoder(&seen), fakeGeneration("fake-model"), |
| 52 | + BuildOptions{BatchSize: 32}) |
| 53 | + require.NoError(t, err) |
| 54 | + assert.Equal(t, []string{"real content"}, seen, |
| 55 | + "whitespace-only content must never reach the embeddings endpoint") |
| 56 | + assert.True(t, result.Activated, |
| 57 | + "a stamped-without-vectors blank document still counts as covered") |
| 58 | + |
| 59 | + var stamps int |
| 60 | + require.NoError(t, ix.db.QueryRow( |
| 61 | + `SELECT COUNT(*) FROM message_vectors_stamps`).Scan(&stamps)) |
| 62 | + assert.Equal(t, 2, stamps, "both documents are stamped") |
| 63 | + |
| 64 | + var chunks int |
| 65 | + require.NoError(t, ix.db.QueryRow( |
| 66 | + `SELECT COUNT(*) FROM message_vectors_chunks`).Scan(&chunks)) |
| 67 | + assert.Equal(t, 1, chunks, "only the non-blank document has a chunk") |
| 68 | +} |
| 69 | + |
| 70 | +// TestEmptyEmbeddingInputIsPermanent pins the structured signal that replaced |
| 71 | +// sniffing provider error bodies for whitespace wording: kit refuses a blank |
| 72 | +// chunk before any HTTP call, and that refusal is a permanent per-document |
| 73 | +// rejection, so a fill stamp-skips it instead of retrying it forever. |
| 74 | +func TestEmptyEmbeddingInputIsPermanent(t *testing.T) { |
| 75 | + err := fmt.Errorf("encode chunk 0: %w", kitvec.ErrEmptyEmbeddingInput) |
| 76 | + assert.True(t, isPermanentEncodeError(err)) |
| 77 | +} |
| 78 | + |
| 79 | +// TestBuildSkipsPermanentlyRejectedDocumentSharingABatch is the cross-document |
| 80 | +// batching counterpart to TestBuildSkipsPermanentlyRejectedDocumentAndContinues. |
| 81 | +// A configured batch_size (production always sets one) packs chunks from |
| 82 | +// several documents into one encode call, so a permanent rejection arrives |
| 83 | +// with no attribution. The build must still isolate the offending document and |
| 84 | +// skip only it; aborting would wedge every later build at the same document. |
| 85 | +func TestBuildSkipsPermanentlyRejectedDocumentSharingABatch(t *testing.T) { |
| 86 | + ix := openTestIndex(t) |
| 87 | + ctx := context.Background() |
| 88 | + src := &fakeUnitSource{rows: []fakeUnit{ |
| 89 | + {unit: userDoc("s1", "", 0, "one"), endedAt: "2024-01-01T00:00:00Z"}, |
| 90 | + {unit: userDoc("s1", "", 1, "poison"), endedAt: "2024-01-01T00:00:01Z"}, |
| 91 | + {unit: userDoc("s1", "", 2, "three"), endedAt: "2024-01-01T00:00:02Z"}, |
| 92 | + }} |
| 93 | + |
| 94 | + rejectPoison := func(_ context.Context, texts []string) ([][]float32, error) { |
| 95 | + if slices.Contains(texts, "poison") { |
| 96 | + return nil, &HTTPStatusError{ |
| 97 | + Status: 400, Body: "input exceeds maximum context length", |
| 98 | + } |
| 99 | + } |
| 100 | + out := make([][]float32, len(texts)) |
| 101 | + for i := range texts { |
| 102 | + out[i] = []float32{1, 0, 0} |
| 103 | + } |
| 104 | + return out, nil |
| 105 | + } |
| 106 | + |
| 107 | + result, err := ix.Build(ctx, src, rejectPoison, fakeGeneration("fake-model"), |
| 108 | + BuildOptions{BatchSize: 32}) |
| 109 | + require.NoError(t, err, |
| 110 | + "one poison document in a shared batch must not abort the whole build") |
| 111 | + assert.Equal(t, 2, result.Fill.Documents, "the two good documents still embed") |
| 112 | + assert.Equal(t, 1, result.Fill.Skipped, "only the poison document is skipped") |
| 113 | + assert.True(t, result.Activated) |
| 114 | +} |
| 115 | + |
| 116 | +// TestBuildTransientBatchErrorStillAborts guards the other side of batch |
| 117 | +// isolation: a 5xx applies to the call, not to one input, so it must abort |
| 118 | +// without probing each document slice separately. |
| 119 | +func TestBuildTransientBatchErrorStillAborts(t *testing.T) { |
| 120 | + ix := openTestIndex(t) |
| 121 | + ctx := context.Background() |
| 122 | + src := &fakeUnitSource{rows: []fakeUnit{ |
| 123 | + {unit: userDoc("s1", "", 0, "one"), endedAt: "2024-01-01T00:00:00Z"}, |
| 124 | + {unit: userDoc("s1", "", 1, "two"), endedAt: "2024-01-01T00:00:01Z"}, |
| 125 | + }} |
| 126 | + |
| 127 | + var calls int |
| 128 | + failing := func(_ context.Context, _ []string) ([][]float32, error) { |
| 129 | + calls++ |
| 130 | + return nil, &HTTPStatusError{Status: 503, Body: "upstream unavailable"} |
| 131 | + } |
| 132 | + |
| 133 | + result, err := ix.Build(ctx, src, failing, fakeGeneration("fake-model"), |
| 134 | + BuildOptions{BatchSize: 32}) |
| 135 | + require.Error(t, err) |
| 136 | + assert.Zero(t, result.Fill.Skipped, "a transient failure must never skip-stamp") |
| 137 | + assert.Equal(t, 1, calls, "a transient failure must not trigger per-document probes") |
| 138 | +} |
| 139 | + |
| 140 | +// TestChunkSnippetResolvesIndexAcrossADroppedWindow covers a search hit on a |
| 141 | +// document with a blank window: its chunk indexes have a gap, so resolving a |
| 142 | +// snippet by slice position would return the wrong chunk's text (or none). |
| 143 | +func TestChunkSnippetResolvesIndexAcrossADroppedWindow(t *testing.T) { |
| 144 | + split := kitvec.SplitOptions{MaxRunes: 10} |
| 145 | + content := strings.Repeat("a", 10) + strings.Repeat(" ", 10) + strings.Repeat("b", 10) |
| 146 | + chunks := kitvec.Split(content, split) |
| 147 | + require.Len(t, chunks, 2, "the all-whitespace middle window is dropped") |
| 148 | + require.Equal(t, []int{0, 2}, []int{chunks[0].Index, chunks[1].Index}, |
| 149 | + "the surviving chunks keep their window numbers") |
| 150 | + |
| 151 | + assert.Equal(t, strings.Repeat("a", 10), chunkSnippet(content, 0, split)) |
| 152 | + assert.Equal(t, strings.Repeat("b", 10), chunkSnippet(content, 2, split), |
| 153 | + "the second stored chunk resolves by its index, not its slice position") |
| 154 | + assert.Empty(t, chunkSnippet(content, 1, split), |
| 155 | + "the dropped window has no snippet") |
| 156 | +} |
| 157 | + |
| 158 | +// TestRepairKeepsDocumentWithADroppedWindow covers the repair scan's view of |
| 159 | +// the same document: chunk indexes 0 and 2 are exactly what Split asks for, so |
| 160 | +// repair must leave the document alone. Treating the gap as a missing chunk |
| 161 | +// would re-embed the document on every repair run, forever. |
| 162 | +func TestRepairKeepsDocumentWithADroppedWindow(t *testing.T) { |
| 163 | + ix := openTestIndex(t) |
| 164 | + ctx := context.Background() |
| 165 | + src := &fakeUnitSource{rows: []fakeUnit{ |
| 166 | + {unit: userDoc("s1", "", 0, blankWindowContent()), endedAt: "2024-01-01T00:00:00Z"}, |
| 167 | + }} |
| 168 | + gen := fakeGeneration("fake-model") |
| 169 | + |
| 170 | + var seen []string |
| 171 | + built, err := ix.Build(ctx, src, recordingEncoder(&seen), gen, BuildOptions{BatchSize: 32}) |
| 172 | + require.NoError(t, err) |
| 173 | + require.Equal(t, 2, built.Fill.Chunks) |
| 174 | + |
| 175 | + repaired, err := ix.Build(ctx, src, recordingEncoder(&seen), gen, |
| 176 | + BuildOptions{BatchSize: 32, RepairInvalid: true}) |
| 177 | + require.NoError(t, err) |
| 178 | + assert.Zero(t, repaired.Repair.Documents, |
| 179 | + "a document whose chunk indexes skip a blank window is healthy") |
| 180 | +} |
0 commit comments