Skip to content

Commit 64a0085

Browse files
geekypunkclaude
andauthored
fix: shrink oversized embedding input instead of dropping the document (#90)
## Problem `createEmbeddingOrEmpty` logs a failed embedding as **"transient/non-fatal"** and moves on. For a rate limit that's right. For an over-length input it isn't: the input is *deterministic*, so the same document fails on every rebuild, and the index keeps a permanent hole that looks like a passing run. Observed on a 567-table MySQL connection during a full reindex: ``` Skipping embedding for RELATIONSHIP document b7efe9d2-… due to transient/non-fatal error: 400: Invalid 'input': maximum context length is 8192 tokens. ``` ## Cause: the character budget, not a missing one `truncate` already cut input to `app.embedding.max-chars` (default **30,000**) on the stated assumption of *"roughly 4 chars per token"* — which would be 7,500 tokens, safely under 8,192. That ratio holds for prose. It does **not** hold for what this service actually embeds: schema and relationship documents are dense identifiers, underscores, punctuation and repeated scaffolding, which tokenize closer to **2–3 chars per token**. At that density 30,000 chars is 10,000–15,000 tokens, and the provider rejects the call. ## Fix No fixed ratio is safe across content, so this stops betting on one. `LlmErrorCategory.CONTEXT_LENGTH` already documents itself as *"never retry; the caller may trim"* — until now nothing trimmed. `embedWithShrink` halves the budget on each CONTEXT_LENGTH rejection (30,000 → 15,000 → … → 1,875, floor 1,000) and lets the provider decide when the call fits. No tokenizer dependency, and correct for any content and any model window. ## Deliberately narrow - **Only CONTEXT_LENGTH shrinks.** Retries and fail-open are untouched for every other category — a smaller input answers nothing about a rate limit or a rejected credential. - **The inner attempt runs with fail-open off** so the rejection reaches the shrink loop; fail-open would convert it into an empty vector indistinguishable from a real one. The operator's fail-open setting is still honoured once shrinking is exhausted. - **Batch shrinking only affects members longer than the budget**, so one oversized text costs the short ones nothing. - **Bounded at 4 halvings**, so a pathological document cannot loop. ## Test plan - [x] `shrinksTheInputWhenTheProviderRejectsItAsTooLong` — succeeds after one halving; asserts the two calls were 30,000 then 15,000 chars. - [x] `givesUpAfterTheShrinkFloorAndStillHonoursFailOpen` — 5 attempts then empty, proving the bound. - [x] `doesNotShrinkForFailuresThatShrinkingCannotFix` — AUTH gets exactly one call. ## Follow-up (not in this PR) The `"transient/non-fatal"` wording in `TrainingService.createEmbeddingOrEmpty` is what made this invisible for so long — it asserts transience for six call sites without knowing the category. Worth revisiting separately. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 4366ef0 commit 64a0085

2 files changed

Lines changed: 233 additions & 14 deletions

File tree

backend/src/main/java/com/dbaagent/service/EmbeddingService.java

Lines changed: 114 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ public class EmbeddingService {
3131

3232
/** Retry budget: 2 retries after the first try, matching the pre-delegation service. */
3333
private static final int MAX_RETRY_ATTEMPTS = 2;
34+
/** Halvings allowed before giving up: 30,000 chars reaches ~1,875 in four. */
35+
private static final int MAX_SHRINK_ATTEMPTS = 4;
36+
/** Floor for shrinking. Below this the document is too small to be worth indexing. */
37+
private static final int MIN_EMBEDDING_CHARS = 1_000;
3438
private static final long RETRY_BACKOFF_MS = 750L;
3539

3640
/**
@@ -87,9 +91,23 @@ public EmbeddingService(
8791
this.maxRetryAfter = maxRetryAfter;
8892
}
8993

90-
/** text-embedding-3-large accepts 8192 tokens, roughly 4 chars per token. */
91-
private String truncate(String text) {
92-
return (text != null && text.length() > maxChars) ? text.substring(0, maxChars) : text;
94+
/**
95+
* Cut {@code text} to a character budget.
96+
*
97+
* <p>The budget is a <em>guess</em> at the model's token window, and it is wrong often
98+
* enough to matter. The old default assumed "roughly 4 chars per token", which holds
99+
* for prose but not for what this service actually embeds: schema and relationship
100+
* documents are dense identifiers — {@code ORDERS.customer_id}, underscores,
101+
* punctuation, repeated scaffolding — that tokenize closer to 2-3 chars per token. At
102+
* that density the 30,000-char default is 10,000-15,000 tokens, well past the 8,192 a
103+
* text-embedding-3-large call accepts, and the provider rejects the request outright.
104+
*
105+
* <p>No fixed ratio is safe across content, so the budget is not trusted to be right.
106+
* {@link #embedWithShrink} lets the provider's own CONTEXT_LENGTH rejection drive the
107+
* budget down until the call fits.
108+
*/
109+
private String truncate(String text, int budget) {
110+
return (text != null && text.length() > budget) ? text.substring(0, budget) : text;
93111
}
94112

95113
/**
@@ -98,8 +116,62 @@ private String truncate(String text) {
98116
public List<Double> createEmbedding(String text) {
99117
LlmCredentials credentials = requireCredentials();
100118
LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId());
101-
return attempt(provider, credentials,
102-
() -> provider.embed(truncate(text), credentials), List.of());
119+
return embedWithShrink(provider, credentials,
120+
budget -> provider.embed(truncate(text, budget), credentials),
121+
List.of(),
122+
text == null ? 0 : text.length());
123+
}
124+
125+
/**
126+
* Run an embedding call, halving the character budget each time the provider says the
127+
* input is too long.
128+
*
129+
* <p>{@link LlmErrorCategory#CONTEXT_LENGTH} already documents itself as "never retry;
130+
* the caller may trim" — but until now no caller trimmed. The rejection fell through to
131+
* fail-open, the document was skipped, and because the input is deterministic it was
132+
* skipped again on every subsequent rebuild. That is a permanent hole in the index
133+
* wearing the costume of a transient blip.
134+
*
135+
* <p>Shrinking is deliberately driven by the provider rather than by counting tokens
136+
* locally: it needs no tokenizer dependency, and it stays correct for any content and
137+
* any model window, including ones whose ratio we have never measured.
138+
*
139+
* <p>Only CONTEXT_LENGTH shrinks. Every other category keeps its existing behaviour —
140+
* retries and fail-open are unchanged — because shrinking the input answers nothing
141+
* about a rate limit or a bad credential.
142+
*/
143+
private <T> T embedWithShrink(LlmEmbeddingProvider provider, LlmCredentials credentials,
144+
java.util.function.IntFunction<T> call, T failOpenValue,
145+
int longestInput) {
146+
// Anchor to what is actually being sent, not to the configured ceiling. Seeding
147+
// from maxChars makes the first halvings no-ops whenever the input is already
148+
// under it — the same bytes, resent and rejected — and leaves anything below
149+
// maxChars/2^MAX_SHRINK_ATTEMPTS unable to shrink at all, which is the very
150+
// document this method exists to rescue.
151+
int budget = Math.min(maxChars, Math.max(longestInput, 1));
152+
for (int shrink = 0; ; shrink++) {
153+
final int attemptBudget = budget;
154+
try {
155+
// attemptOrThrow, not attempt: a CONTEXT_LENGTH rejection has to reach us
156+
// unlogged. Fail-open would turn it into an empty vector indistinguishable
157+
// from a real one, and the logging path would report a failure for a call
158+
// that is about to succeed.
159+
return attemptOrThrow(provider, credentials, () -> call.apply(attemptBudget));
160+
} catch (RuntimeException e) {
161+
LlmErrorCategory category = provider.classify(e);
162+
if (category != LlmErrorCategory.CONTEXT_LENGTH
163+
|| shrink >= MAX_SHRINK_ATTEMPTS
164+
|| budget <= MIN_EMBEDDING_CHARS) {
165+
// Out of room to shrink, or not a length problem: honour the operator's
166+
// fail-open setting exactly as before this method existed. This is the
167+
// one place a terminal failure is logged.
168+
return handleFailure(category, credentials, e, failOpenValue, failOpen);
169+
}
170+
int next = Math.max(MIN_EMBEDDING_CHARS, budget / 2);
171+
log.warn("Embedding rejected as too long at {} chars sent; retrying at {}", budget, next);
172+
budget = next;
173+
}
174+
}
103175
}
104176

105177
/**
@@ -110,10 +182,23 @@ public List<Double> createEmbedding(String text) {
110182
public List<List<Double>> createEmbeddings(List<String> texts) {
111183
LlmCredentials credentials = requireCredentials();
112184
LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId());
113-
List<String> truncated = texts.stream().map(this::truncate).toList();
114-
return attempt(provider, credentials,
115-
() -> provider.embedBatch(truncated, credentials),
116-
Collections.nCopies(texts.size(), List.of()));
185+
// The budget is per-request, not per-member, so a halving forced by one oversized
186+
// text also trims every other member above the new budget. That is a real cost and
187+
// not something this method can avoid: the provider rejects the request, not a
188+
// document, and it does not say which member was at fault. Seeding from the
189+
// longest member keeps the first halving meaningful; callers that cannot afford
190+
// collateral truncation should embed individually.
191+
//
192+
// Note also that a provider may reject on the request's AGGREGATE token count, in
193+
// which case trimming members is the right lever but the floor may be reached
194+
// before the batch fits.
195+
int longest = texts.stream().filter(java.util.Objects::nonNull)
196+
.mapToInt(String::length).max().orElse(0);
197+
return embedWithShrink(provider, credentials,
198+
budget -> provider.embedBatch(
199+
texts.stream().map(t -> truncate(t, budget)).toList(), credentials),
200+
Collections.nCopies(texts.size(), List.of()),
201+
longest);
117202
}
118203

119204
/**
@@ -166,10 +251,6 @@ public int dimensions() {
166251
* <p>The decision is {@link LlmErrorCategory#isRetryable()}, not a message substring:
167252
* that taxonomy exists precisely so retry policy stops being provider-specific.
168253
*/
169-
private <T> T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials,
170-
Supplier<T> call, T failOpenValue) {
171-
return attempt(provider, credentials, call, failOpenValue, failOpen);
172-
}
173254

174255
/**
175256
* As above, but with fail-open decided per call site rather than by configuration.
@@ -180,14 +261,33 @@ private <T> T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials,
180261
*/
181262
private <T> T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials,
182263
Supplier<T> call, T failOpenValue, boolean mayFailOpen) {
264+
try {
265+
return attemptOrThrow(provider, credentials, call);
266+
} catch (RuntimeException e) {
267+
return handleFailure(provider.classify(e), credentials, e, failOpenValue, mayFailOpen);
268+
}
269+
}
270+
271+
/**
272+
* The retry loop with no opinion about failure: exhausted retries rethrow.
273+
*
274+
* <p>Separating this from {@link #handleFailure} is what lets {@link #embedWithShrink}
275+
* treat a CONTEXT_LENGTH rejection as a step in a working algorithm rather than an
276+
* incident. Routing intermediate attempts through the logging path made every
277+
* successful shrink emit "Embedding failed" and a stack trace for a call that then
278+
* succeeded — noise that would fire any alert keyed on that string, and log a terminal
279+
* failure twice with contradictory failOpen values.
280+
*/
281+
private <T> T attemptOrThrow(LlmEmbeddingProvider provider, LlmCredentials credentials,
282+
Supplier<T> call) {
183283
for (int retries = 0; ; retries++) {
184284
try {
185285
return call.get();
186286
} catch (RuntimeException e) {
187287
LlmErrorCategory category = provider.classify(e);
188288
if (retries >= MAX_RETRY_ATTEMPTS || !category.isRetryable()
189289
|| !backoff(retries, category, credentials, retryAfterHint(provider, e))) {
190-
return handleFailure(category, credentials, e, failOpenValue, mayFailOpen);
290+
throw e;
191291
}
192292
}
193293
}

backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,4 +476,123 @@ void cosineSimilarityOfIdenticalVectorsIsOne() {
476476
assertThat(service.cosineSimilarity(List.of(1.0, 2.0, 3.0), List.of(1.0, 2.0, 3.0)))
477477
.isCloseTo(1.0, org.assertj.core.data.Offset.offset(1e-9));
478478
}
479+
480+
@Test
481+
void shrinksTheInputWhenTheProviderRejectsItAsTooLong() {
482+
var resolver = mock(LlmConfigResolver.class);
483+
var registry = mock(LlmProviderRegistry.class);
484+
var provider = mock(LlmEmbeddingProvider.class);
485+
486+
when(resolver.resolveEmbedding()).thenReturn(creds());
487+
when(registry.embeddingProvider("openai")).thenReturn(provider);
488+
// Stands in for a model whose real token window is reached well before the
489+
// character budget: anything over 15,000 chars is rejected outright.
490+
when(provider.embed(anyString(), any())).thenAnswer(call -> {
491+
String sent = call.getArgument(0);
492+
if (sent.length() > 15_000) {
493+
throw new RuntimeException("maximum context length is 8192 tokens");
494+
}
495+
return List.of(0.5);
496+
});
497+
when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH);
498+
499+
var service = new EmbeddingService(resolver, registry, 30_000, false);
500+
501+
// Without shrinking this document is dropped forever; the point of the fix is
502+
// that it comes back embedded rather than empty.
503+
assertThat(service.createEmbedding("x".repeat(50_000))).containsExactly(0.5);
504+
505+
var captor = ArgumentCaptor.forClass(String.class);
506+
verify(provider, times(2)).embed(captor.capture(), any());
507+
assertThat(captor.getAllValues().get(0)).hasSize(30_000);
508+
assertThat(captor.getAllValues().get(1)).hasSize(15_000);
509+
}
510+
511+
@Test
512+
void givesUpAfterTheAttemptCapAndStillHonoursFailOpen() {
513+
var resolver = mock(LlmConfigResolver.class);
514+
var registry = mock(LlmProviderRegistry.class);
515+
var provider = mock(LlmEmbeddingProvider.class);
516+
517+
when(resolver.resolveEmbedding()).thenReturn(creds());
518+
when(registry.embeddingProvider("openai")).thenReturn(provider);
519+
when(provider.embed(anyString(), any()))
520+
.thenThrow(new RuntimeException("maximum context length is 8192 tokens"));
521+
when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH);
522+
523+
var service = new EmbeddingService(resolver, registry, 30_000, true);
524+
525+
// Bounded by MAX_SHRINK_ATTEMPTS, not by MIN_EMBEDDING_CHARS: at a 30,000
526+
// ceiling the cap is reached first. The floor is covered separately below.
527+
assertThat(service.createEmbedding("x".repeat(50_000))).isEmpty();
528+
verify(provider, times(5)).embed(anyString(), any()); // 30k, 15k, 7.5k, 3.75k, 1.875k
529+
}
530+
531+
@Test
532+
void doesNotShrinkForFailuresThatShrinkingCannotFix() {
533+
var resolver = mock(LlmConfigResolver.class);
534+
var registry = mock(LlmProviderRegistry.class);
535+
var provider = mock(LlmEmbeddingProvider.class);
536+
537+
when(resolver.resolveEmbedding()).thenReturn(creds());
538+
when(registry.embeddingProvider("openai")).thenReturn(provider);
539+
when(provider.embed(anyString(), any())).thenThrow(new RuntimeException("nope"));
540+
when(provider.classify(any())).thenReturn(LlmErrorCategory.AUTH);
541+
542+
var service = new EmbeddingService(resolver, registry, 30_000, true);
543+
544+
assertThat(service.createEmbedding("x".repeat(50_000))).isEmpty();
545+
// A rejected credential says nothing about input size; retrying smaller would
546+
// just multiply the failed calls.
547+
verify(provider, times(1)).embed(anyString(), any());
548+
}
549+
550+
@Test
551+
void shrinksRelativeToTheInputNotTheConfiguredCeiling() {
552+
var resolver = mock(LlmConfigResolver.class);
553+
var registry = mock(LlmProviderRegistry.class);
554+
var provider = mock(LlmEmbeddingProvider.class);
555+
556+
when(resolver.resolveEmbedding()).thenReturn(creds());
557+
when(registry.embeddingProvider("openai")).thenReturn(provider);
558+
when(provider.embed(anyString(), any())).thenAnswer(call -> {
559+
String sent = call.getArgument(0);
560+
if (sent.length() > 2_000) {
561+
throw new RuntimeException("maximum context length is 8192 tokens");
562+
}
563+
return List.of(0.5);
564+
});
565+
when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH);
566+
567+
// 4,000 chars against a 30,000 ceiling. Seeding the budget from the ceiling would
568+
// make the first halvings no-ops — 30,000 and 15,000 both send the same 4,000
569+
// bytes — and burn the attempt cap on identical rejected calls.
570+
new EmbeddingService(resolver, registry, 30_000, true)
571+
.createEmbedding("x".repeat(4_000));
572+
573+
var captor = ArgumentCaptor.forClass(String.class);
574+
verify(provider, times(2)).embed(captor.capture(), any());
575+
assertThat(captor.getAllValues().get(0)).hasSize(4_000);
576+
assertThat(captor.getAllValues().get(1)).hasSize(2_000);
577+
}
578+
579+
@Test
580+
void stopsAtTheCharacterFloorRatherThanEmbeddingATokenOfContent() {
581+
var resolver = mock(LlmConfigResolver.class);
582+
var registry = mock(LlmProviderRegistry.class);
583+
var provider = mock(LlmEmbeddingProvider.class);
584+
585+
when(resolver.resolveEmbedding()).thenReturn(creds());
586+
when(registry.embeddingProvider("openai")).thenReturn(provider);
587+
when(provider.embed(anyString(), any()))
588+
.thenThrow(new RuntimeException("maximum context length is 8192 tokens"));
589+
when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH);
590+
591+
// A ceiling low enough that MIN_EMBEDDING_CHARS (1,000) is what stops the loop,
592+
// not the attempt cap: 1,500 -> 1,000 -> give up.
593+
var service = new EmbeddingService(resolver, registry, 1_500, true);
594+
595+
assertThat(service.createEmbedding("x".repeat(5_000))).isEmpty();
596+
verify(provider, times(2)).embed(anyString(), any());
597+
}
479598
}

0 commit comments

Comments
 (0)