From f93f4831331d339c03c5d4f2d3aeaf7ca399f4c2 Mon Sep 17 00:00:00 2001 From: Sasank Talasila Date: Fri, 28 Aug 2026 16:53:04 -0500 Subject: [PATCH 1/2] fix: shrink oversized embedding input instead of dropping the document MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `createEmbeddingOrEmpty` logs a failed embedding as "transient/non-fatal" and moves on. For a rate limit that is right. For an over-length input it is not: 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. The cause is 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. That ratio holds for prose. It does not hold for what this service 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. 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. That needs no tokenizer dependency and stays correct for any content and any model window. Deliberately narrow: - Only CONTEXT_LENGTH shrinks. Retries and fail-open are untouched for every other category, since 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. - Shrinking a batch 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. Tests cover all three paths: shrink-then-succeed, exhaust-then-fail-open, and no-shrink for a category shrinking cannot fix. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 5 (1M context) --- .../dbaagent/service/EmbeddingService.java | 80 +++++++++++++++++-- .../service/EmbeddingServiceTest.java | 69 ++++++++++++++++ 2 files changed, 141 insertions(+), 8 deletions(-) diff --git a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java index 835094b..afbcdc8 100644 --- a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java +++ b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java @@ -31,6 +31,10 @@ public class EmbeddingService { /** Retry budget: 2 retries after the first try, matching the pre-delegation service. */ private static final int MAX_RETRY_ATTEMPTS = 2; + /** Halvings allowed before giving up: 30,000 chars reaches ~1,875 in four. */ + private static final int MAX_SHRINK_ATTEMPTS = 4; + /** Floor for shrinking. Below this the document is too small to be worth indexing. */ + private static final int MIN_EMBEDDING_CHARS = 1_000; private static final long RETRY_BACKOFF_MS = 750L; /** @@ -87,9 +91,23 @@ public EmbeddingService( this.maxRetryAfter = maxRetryAfter; } - /** text-embedding-3-large accepts 8192 tokens, roughly 4 chars per token. */ - private String truncate(String text) { - return (text != null && text.length() > maxChars) ? text.substring(0, maxChars) : text; + /** + * Cut {@code text} to a character budget. + * + *

The budget is a guess at the model's token window, and it is wrong often + * enough to matter. The old default assumed "roughly 4 chars per token", which holds + * for prose but not for what this service actually embeds: schema and relationship + * documents are dense identifiers — {@code ORDERS.customer_id}, underscores, + * punctuation, repeated scaffolding — that tokenize closer to 2-3 chars per token. At + * that density the 30,000-char default is 10,000-15,000 tokens, well past the 8,192 a + * text-embedding-3-large call accepts, and the provider rejects the request outright. + * + *

No fixed ratio is safe across content, so the budget is not trusted to be right. + * {@link #embedWithShrink} lets the provider's own CONTEXT_LENGTH rejection drive the + * budget down until the call fits. + */ + private String truncate(String text, int budget) { + return (text != null && text.length() > budget) ? text.substring(0, budget) : text; } /** @@ -98,8 +116,52 @@ private String truncate(String text) { public List createEmbedding(String text) { LlmCredentials credentials = requireCredentials(); LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId()); - return attempt(provider, credentials, - () -> provider.embed(truncate(text), credentials), List.of()); + return embedWithShrink(provider, credentials, + budget -> provider.embed(truncate(text, budget), credentials), + List.of()); + } + + /** + * Run an embedding call, halving the character budget each time the provider says the + * input is too long. + * + *

{@link LlmErrorCategory#CONTEXT_LENGTH} already documents itself as "never retry; + * the caller may trim" — but until now no caller trimmed. The rejection fell through to + * fail-open, the document was skipped, and because the input is deterministic it was + * skipped again on every subsequent rebuild. That is a permanent hole in the index + * wearing the costume of a transient blip. + * + *

Shrinking is deliberately driven by the provider rather than by counting tokens + * locally: it needs no tokenizer dependency, and it stays correct for any content and + * any model window, including ones whose ratio we have never measured. + * + *

Only CONTEXT_LENGTH shrinks. Every other category keeps its existing behaviour — + * retries and fail-open are unchanged — because shrinking the input answers nothing + * about a rate limit or a bad credential. + */ + private T embedWithShrink(LlmEmbeddingProvider provider, LlmCredentials credentials, + java.util.function.IntFunction call, T failOpenValue) { + int budget = maxChars; + for (int shrink = 0; ; shrink++) { + final int attemptBudget = budget; + try { + // mayFailOpen=false: a CONTEXT_LENGTH rejection has to reach us. Fail-open + // would turn it into an empty vector indistinguishable from a real one. + return attempt(provider, credentials, () -> call.apply(attemptBudget), failOpenValue, false); + } catch (RuntimeException e) { + LlmErrorCategory category = provider.classify(e); + if (category != LlmErrorCategory.CONTEXT_LENGTH + || shrink >= MAX_SHRINK_ATTEMPTS + || budget <= MIN_EMBEDDING_CHARS) { + // Out of room to shrink, or not a length problem: honour the operator's + // fail-open setting exactly as before this method existed. + return handleFailure(category, credentials, e, failOpenValue, failOpen); + } + int next = Math.max(MIN_EMBEDDING_CHARS, budget / 2); + log.warn("Embedding input rejected as too long at {} chars; retrying at {}", budget, next); + budget = next; + } + } } /** @@ -110,9 +172,11 @@ public List createEmbedding(String text) { public List> createEmbeddings(List texts) { LlmCredentials credentials = requireCredentials(); LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId()); - List truncated = texts.stream().map(this::truncate).toList(); - return attempt(provider, credentials, - () -> provider.embedBatch(truncated, credentials), + // Shrinking the budget only touches texts longer than it, so one oversized member + // of a batch cannot cost the short ones any content. + return embedWithShrink(provider, credentials, + budget -> provider.embedBatch( + texts.stream().map(t -> truncate(t, budget)).toList(), credentials), Collections.nCopies(texts.size(), List.of())); } diff --git a/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java b/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java index 56fadea..b60eec6 100644 --- a/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java @@ -476,4 +476,73 @@ void cosineSimilarityOfIdenticalVectorsIsOne() { assertThat(service.cosineSimilarity(List.of(1.0, 2.0, 3.0), List.of(1.0, 2.0, 3.0))) .isCloseTo(1.0, org.assertj.core.data.Offset.offset(1e-9)); } + + @Test + void shrinksTheInputWhenTheProviderRejectsItAsTooLong() { + var resolver = mock(LlmConfigResolver.class); + var registry = mock(LlmProviderRegistry.class); + var provider = mock(LlmEmbeddingProvider.class); + + when(resolver.resolveEmbedding()).thenReturn(creds()); + when(registry.embeddingProvider("openai")).thenReturn(provider); + // Stands in for a model whose real token window is reached well before the + // character budget: anything over 15,000 chars is rejected outright. + when(provider.embed(anyString(), any())).thenAnswer(call -> { + String sent = call.getArgument(0); + if (sent.length() > 15_000) { + throw new RuntimeException("maximum context length is 8192 tokens"); + } + return List.of(0.5); + }); + when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH); + + var service = new EmbeddingService(resolver, registry, 30_000, false); + + // Without shrinking this document is dropped forever; the point of the fix is + // that it comes back embedded rather than empty. + assertThat(service.createEmbedding("x".repeat(50_000))).containsExactly(0.5); + + var captor = ArgumentCaptor.forClass(String.class); + verify(provider, times(2)).embed(captor.capture(), any()); + assertThat(captor.getAllValues().get(0)).hasSize(30_000); + assertThat(captor.getAllValues().get(1)).hasSize(15_000); + } + + @Test + void givesUpAfterTheShrinkFloorAndStillHonoursFailOpen() { + var resolver = mock(LlmConfigResolver.class); + var registry = mock(LlmProviderRegistry.class); + var provider = mock(LlmEmbeddingProvider.class); + + when(resolver.resolveEmbedding()).thenReturn(creds()); + when(registry.embeddingProvider("openai")).thenReturn(provider); + when(provider.embed(anyString(), any())) + .thenThrow(new RuntimeException("maximum context length is 8192 tokens")); + when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH); + + var service = new EmbeddingService(resolver, registry, 30_000, true); + + // Shrinking is bounded, so a pathological document cannot loop forever. + assertThat(service.createEmbedding("x".repeat(50_000))).isEmpty(); + verify(provider, times(5)).embed(anyString(), any()); // 30k, 15k, 7.5k, 3.75k, 1.875k + } + + @Test + void doesNotShrinkForFailuresThatShrinkingCannotFix() { + var resolver = mock(LlmConfigResolver.class); + var registry = mock(LlmProviderRegistry.class); + var provider = mock(LlmEmbeddingProvider.class); + + when(resolver.resolveEmbedding()).thenReturn(creds()); + when(registry.embeddingProvider("openai")).thenReturn(provider); + when(provider.embed(anyString(), any())).thenThrow(new RuntimeException("nope")); + when(provider.classify(any())).thenReturn(LlmErrorCategory.AUTH); + + var service = new EmbeddingService(resolver, registry, 30_000, true); + + assertThat(service.createEmbedding("x".repeat(50_000))).isEmpty(); + // A rejected credential says nothing about input size; retrying smaller would + // just multiply the failed calls. + verify(provider, times(1)).embed(anyString(), any()); + } } From 2ed8478064375871044b184962d2eb008019e45d Mon Sep 17 00:00:00 2001 From: Sasank Talasila Date: Fri, 28 Aug 2026 18:02:40 -0500 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20address=20review=20=E2=80=94=20ancho?= =?UTF-8?q?r=20the=20shrink=20budget=20to=20the=20payload,=20not=20the=20c?= =?UTF-8?q?eiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six issues from review, one of which partly defeated the original fix. 1. The budget was seeded from `maxChars` rather than from the text being sent, so `truncate` (which only cuts when length > budget) made the first halvings byte-identical resends. A document below maxChars/2^4 — 1,875 chars at the default — could never shrink at all: five identical rejected calls and then the same permanent index hole this change exists to close, at 5x the cost. The budget now starts at min(maxChars, longest input). 2. The batch comment claimed one oversized member "cannot cost the short ones any content". Untrue: the budget is per-request, so a halving forced by one text also trims every other member above the new budget. The comment now says so, and seeds from the longest member. Also notes that a provider may reject on the request's aggregate token count, where trimming members is the right lever but the floor may arrive before the batch fits. 3. Intermediate attempts ran through `attempt(..., mayFailOpen=false)`, whose `handleFailure` logs before it rethrows — so every successful shrink emitted "Embedding failed" with a stack trace for a call that then succeeded, and a terminal failure logged twice with contradictory failOpen values. The retry loop is now `attemptOrThrow`, with logging left to the single terminal `handleFailure`. 4. The 4-arg `attempt` overload became unreachable once both callers moved to `embedWithShrink`; removed. 5. `givesUpAfterTheShrinkFloor...` asserted the attempt cap, not the floor — at a 30,000 ceiling the cap is always reached first. Renamed to say what it tests, and a real floor test added at a 1,500 ceiling. 6. The shrink log printed the budget as though it were the payload size, which would mislead exactly the person debugging (1). New tests: shrinksRelativeToTheInputNotTheConfiguredCeiling (regression for 1, fails on the previous commit) and stopsAtTheCharacterFloorRatherThanEmbedding ATokenOfContent (coverage for 5). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 5 (1M context) --- .../dbaagent/service/EmbeddingService.java | 68 ++++++++++++++----- .../service/EmbeddingServiceTest.java | 54 ++++++++++++++- 2 files changed, 104 insertions(+), 18 deletions(-) diff --git a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java index afbcdc8..a2d1742 100644 --- a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java +++ b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java @@ -118,7 +118,8 @@ public List createEmbedding(String text) { LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId()); return embedWithShrink(provider, credentials, budget -> provider.embed(truncate(text, budget), credentials), - List.of()); + List.of(), + text == null ? 0 : text.length()); } /** @@ -140,25 +141,34 @@ public List createEmbedding(String text) { * about a rate limit or a bad credential. */ private T embedWithShrink(LlmEmbeddingProvider provider, LlmCredentials credentials, - java.util.function.IntFunction call, T failOpenValue) { - int budget = maxChars; + java.util.function.IntFunction call, T failOpenValue, + int longestInput) { + // Anchor to what is actually being sent, not to the configured ceiling. Seeding + // from maxChars makes the first halvings no-ops whenever the input is already + // under it — the same bytes, resent and rejected — and leaves anything below + // maxChars/2^MAX_SHRINK_ATTEMPTS unable to shrink at all, which is the very + // document this method exists to rescue. + int budget = Math.min(maxChars, Math.max(longestInput, 1)); for (int shrink = 0; ; shrink++) { final int attemptBudget = budget; try { - // mayFailOpen=false: a CONTEXT_LENGTH rejection has to reach us. Fail-open - // would turn it into an empty vector indistinguishable from a real one. - return attempt(provider, credentials, () -> call.apply(attemptBudget), failOpenValue, false); + // attemptOrThrow, not attempt: a CONTEXT_LENGTH rejection has to reach us + // unlogged. Fail-open would turn it into an empty vector indistinguishable + // from a real one, and the logging path would report a failure for a call + // that is about to succeed. + return attemptOrThrow(provider, credentials, () -> call.apply(attemptBudget)); } catch (RuntimeException e) { LlmErrorCategory category = provider.classify(e); if (category != LlmErrorCategory.CONTEXT_LENGTH || shrink >= MAX_SHRINK_ATTEMPTS || budget <= MIN_EMBEDDING_CHARS) { // Out of room to shrink, or not a length problem: honour the operator's - // fail-open setting exactly as before this method existed. + // fail-open setting exactly as before this method existed. This is the + // one place a terminal failure is logged. return handleFailure(category, credentials, e, failOpenValue, failOpen); } int next = Math.max(MIN_EMBEDDING_CHARS, budget / 2); - log.warn("Embedding input rejected as too long at {} chars; retrying at {}", budget, next); + log.warn("Embedding rejected as too long at {} chars sent; retrying at {}", budget, next); budget = next; } } @@ -172,12 +182,23 @@ private T embedWithShrink(LlmEmbeddingProvider provider, LlmCredentials cred public List> createEmbeddings(List texts) { LlmCredentials credentials = requireCredentials(); LlmEmbeddingProvider provider = registry.embeddingProvider(credentials.providerId()); - // Shrinking the budget only touches texts longer than it, so one oversized member - // of a batch cannot cost the short ones any content. + // The budget is per-request, not per-member, so a halving forced by one oversized + // text also trims every other member above the new budget. That is a real cost and + // not something this method can avoid: the provider rejects the request, not a + // document, and it does not say which member was at fault. Seeding from the + // longest member keeps the first halving meaningful; callers that cannot afford + // collateral truncation should embed individually. + // + // Note also that a provider may reject on the request's AGGREGATE token count, in + // which case trimming members is the right lever but the floor may be reached + // before the batch fits. + int longest = texts.stream().filter(java.util.Objects::nonNull) + .mapToInt(String::length).max().orElse(0); return embedWithShrink(provider, credentials, budget -> provider.embedBatch( texts.stream().map(t -> truncate(t, budget)).toList(), credentials), - Collections.nCopies(texts.size(), List.of())); + Collections.nCopies(texts.size(), List.of()), + longest); } /** @@ -230,10 +251,6 @@ public int dimensions() { *

The decision is {@link LlmErrorCategory#isRetryable()}, not a message substring: * that taxonomy exists precisely so retry policy stops being provider-specific. */ - private T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials, - Supplier call, T failOpenValue) { - return attempt(provider, credentials, call, failOpenValue, failOpen); - } /** * As above, but with fail-open decided per call site rather than by configuration. @@ -244,6 +261,25 @@ private T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials, */ private T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials, Supplier call, T failOpenValue, boolean mayFailOpen) { + try { + return attemptOrThrow(provider, credentials, call); + } catch (RuntimeException e) { + return handleFailure(provider.classify(e), credentials, e, failOpenValue, mayFailOpen); + } + } + + /** + * The retry loop with no opinion about failure: exhausted retries rethrow. + * + *

Separating this from {@link #handleFailure} is what lets {@link #embedWithShrink} + * treat a CONTEXT_LENGTH rejection as a step in a working algorithm rather than an + * incident. Routing intermediate attempts through the logging path made every + * successful shrink emit "Embedding failed" and a stack trace for a call that then + * succeeded — noise that would fire any alert keyed on that string, and log a terminal + * failure twice with contradictory failOpen values. + */ + private T attemptOrThrow(LlmEmbeddingProvider provider, LlmCredentials credentials, + Supplier call) { for (int retries = 0; ; retries++) { try { return call.get(); @@ -251,7 +287,7 @@ private T attempt(LlmEmbeddingProvider provider, LlmCredentials credentials, LlmErrorCategory category = provider.classify(e); if (retries >= MAX_RETRY_ATTEMPTS || !category.isRetryable() || !backoff(retries, category, credentials, retryAfterHint(provider, e))) { - return handleFailure(category, credentials, e, failOpenValue, mayFailOpen); + throw e; } } } diff --git a/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java b/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java index b60eec6..c51391f 100644 --- a/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java +++ b/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java @@ -509,7 +509,7 @@ void shrinksTheInputWhenTheProviderRejectsItAsTooLong() { } @Test - void givesUpAfterTheShrinkFloorAndStillHonoursFailOpen() { + void givesUpAfterTheAttemptCapAndStillHonoursFailOpen() { var resolver = mock(LlmConfigResolver.class); var registry = mock(LlmProviderRegistry.class); var provider = mock(LlmEmbeddingProvider.class); @@ -522,7 +522,8 @@ void givesUpAfterTheShrinkFloorAndStillHonoursFailOpen() { var service = new EmbeddingService(resolver, registry, 30_000, true); - // Shrinking is bounded, so a pathological document cannot loop forever. + // Bounded by MAX_SHRINK_ATTEMPTS, not by MIN_EMBEDDING_CHARS: at a 30,000 + // ceiling the cap is reached first. The floor is covered separately below. assertThat(service.createEmbedding("x".repeat(50_000))).isEmpty(); verify(provider, times(5)).embed(anyString(), any()); // 30k, 15k, 7.5k, 3.75k, 1.875k } @@ -545,4 +546,53 @@ void doesNotShrinkForFailuresThatShrinkingCannotFix() { // just multiply the failed calls. verify(provider, times(1)).embed(anyString(), any()); } + + @Test + void shrinksRelativeToTheInputNotTheConfiguredCeiling() { + var resolver = mock(LlmConfigResolver.class); + var registry = mock(LlmProviderRegistry.class); + var provider = mock(LlmEmbeddingProvider.class); + + when(resolver.resolveEmbedding()).thenReturn(creds()); + when(registry.embeddingProvider("openai")).thenReturn(provider); + when(provider.embed(anyString(), any())).thenAnswer(call -> { + String sent = call.getArgument(0); + if (sent.length() > 2_000) { + throw new RuntimeException("maximum context length is 8192 tokens"); + } + return List.of(0.5); + }); + when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH); + + // 4,000 chars against a 30,000 ceiling. Seeding the budget from the ceiling would + // make the first halvings no-ops — 30,000 and 15,000 both send the same 4,000 + // bytes — and burn the attempt cap on identical rejected calls. + new EmbeddingService(resolver, registry, 30_000, true) + .createEmbedding("x".repeat(4_000)); + + var captor = ArgumentCaptor.forClass(String.class); + verify(provider, times(2)).embed(captor.capture(), any()); + assertThat(captor.getAllValues().get(0)).hasSize(4_000); + assertThat(captor.getAllValues().get(1)).hasSize(2_000); + } + + @Test + void stopsAtTheCharacterFloorRatherThanEmbeddingATokenOfContent() { + var resolver = mock(LlmConfigResolver.class); + var registry = mock(LlmProviderRegistry.class); + var provider = mock(LlmEmbeddingProvider.class); + + when(resolver.resolveEmbedding()).thenReturn(creds()); + when(registry.embeddingProvider("openai")).thenReturn(provider); + when(provider.embed(anyString(), any())) + .thenThrow(new RuntimeException("maximum context length is 8192 tokens")); + when(provider.classify(any())).thenReturn(LlmErrorCategory.CONTEXT_LENGTH); + + // A ceiling low enough that MIN_EMBEDDING_CHARS (1,000) is what stops the loop, + // not the attempt cap: 1,500 -> 1,000 -> give up. + var service = new EmbeddingService(resolver, registry, 1_500, true); + + assertThat(service.createEmbedding("x".repeat(5_000))).isEmpty(); + verify(provider, times(2)).embed(anyString(), any()); + } }