diff --git a/backend/src/main/java/com/dbaagent/service/EmbeddingService.java b/backend/src/main/java/com/dbaagent/service/EmbeddingService.java
index 835094b..a2d1742 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,62 @@ 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(),
+ text == null ? 0 : text.length());
+ }
+
+ /**
+ * 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 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 {
+ // 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. 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 rejected as too long at {} chars sent; retrying at {}", budget, next);
+ budget = next;
+ }
+ }
}
/**
@@ -110,10 +182,23 @@ 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),
- Collections.nCopies(texts.size(), List.of()));
+ // 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()),
+ longest);
}
/**
@@ -166,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.
@@ -180,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();
@@ -187,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 56fadea..c51391f 100644
--- a/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java
+++ b/backend/src/test/java/com/dbaagent/service/EmbeddingServiceTest.java
@@ -476,4 +476,123 @@ 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 givesUpAfterTheAttemptCapAndStillHonoursFailOpen() {
+ 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);
+
+ // 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
+ }
+
+ @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());
+ }
+
+ @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());
+ }
}