fix(embeddings): keep the batch ceiling exact and warn before truncating

Review round 5. Reverts the safety margin from round 4.

The two review findings were in direct tension: round 4 flagged that a
foreign model's ceiling is measured in tiktoken units, and the margin added
to absorb that error reintroduced the round 3 harm — valid content truncated
below the provider's declared limit.

The margin was the wrong trade. It swapped a loud failure for a silent one:
an undercount surfaces as a provider rejection the caller can see and act on,
while shortening an embedding's input produces a degraded vector that is
indistinguishable from a good one at every layer above it. Silent quality
loss in a retrieval index is the worse outcome, and it is also the harder one
to ever notice.

So the declared ceiling is applied exactly, and truncation is no longer
silent: an input above the limit now logs a warning naming the model, the
limit, and whether the count was approximate. hasApproximateTokenCount
records which models are counted with a foreign tokenizer without being used
to shrink anything.

The tokenizer imprecision itself remains, and cannot be fixed without
per-provider BPE the repo does not have — estimateTokenCount is a
chars-per-token heuristic, and truncation needs a real encode/decode pair to
slice on a token boundary.

Author: mzxchandra

apps/sim/lib/embeddings/catalog.test.ts

@@ -8,8 +8,8 @@ import {
   getEmbeddingModelInfo,
   getKbEligibleModels,
   getModelsForProvider,
+  hasApproximateTokenCount,
   KB_EMBEDDING_DIMENSIONS,
-  resolveBatchTokenCeiling,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 import { EMBEDDING_MODEL_PRICING } from '@/providers/models'
@@ -101,42 +101,30 @@ describe('splitByItemLimit', () => {
     expect(batches.flat()).toEqual(items)
   })
 
-  /**
-   * Batching counts tokens with tiktoken, which only has encodings for OpenAI
-   * models. Every other id falls back to cl100k, so a foreign model's ceiling
-   * would otherwise be enforced in the wrong token units.
-   */
-  describe('resolveBatchTokenCeiling', () => {
-    it('trusts the declared ceiling for tiktoken-native models', () => {
-      for (const [id, info] of Object.entries(EMBEDDING_MODELS)) {
-        if (info.tokenizerProvider !== 'openai') continue
-        expect(resolveBatchTokenCeiling(info), id).toBe(info.maxInputTokens)
-      }
-    })
-
-    it('discounts the ceiling for every model counted with a foreign tokenizer', () => {
-      const foreign = Object.entries(EMBEDDING_MODELS).filter(
-        ([, info]) => info.tokenizerProvider !== 'openai'
-      )
-      // Guards the test itself: the catalog must still contain such models.
-      expect(foreign.length).toBeGreaterThan(0)
-
-      for (const [id, info] of foreign) {
-        const ceiling = resolveBatchTokenCeiling(info)
-        expect(ceiling, id).toBeLessThan(info.maxInputTokens)
-        expect(ceiling, id).toBeGreaterThan(0)
-      }
-    })
-
-    it("leaves Gemini's tight ceiling below its declared 2048", () => {
-      expect(resolveBatchTokenCeiling(getEmbeddingModelInfo('gemini-embedding-001'))).toBe(1638)
-    })
-  })
-
   it("chunks to Cohere's 96-item cap", () => {
     const items = Array.from({ length: 200 }, (_, i) => i)
     const batches = splitByItemLimit(items, 96)
     expect(batches.map((b) => b.length)).toEqual([96, 96, 8])
     expect(batches.flat()).toEqual(items)
   })
 })
+
+/**
+ * Batching counts tokens with tiktoken, which only has encodings for OpenAI
+ * models; every other id falls back to cl100k. This flag records which models
+ * are counted approximately. It must not be used to shrink the ceiling —
+ * truncating below a provider's declared limit drops valid content silently,
+ * which is strictly worse than the visible rejection it would guard against.
+ */
+describe('hasApproximateTokenCount', () => {
+  it('is false only for tiktoken-native models', () => {
+    for (const [id, info] of Object.entries(EMBEDDING_MODELS)) {
+      expect(hasApproximateTokenCount(info), id).toBe(info.tokenizerProvider !== 'openai')
+    }
+  })
+
+  it('covers at least one model, so the flag stays meaningful', () => {
+    const approximate = Object.values(EMBEDDING_MODELS).filter(hasApproximateTokenCount)
+    expect(approximate.length).toBeGreaterThan(0)
+  })
+})

apps/sim/lib/embeddings/catalog.ts

@@ -158,28 +158,18 @@ export function getKbEligibleModels(): string[] {
  * optional caller-requested reduction.
  */
 /**
- * Fraction of a model's declared ceiling used for batching when its tokens are
- * counted with the wrong tokenizer. Chosen to absorb the usual spread between
- * BPE vocabularies on ordinary text; it is a guard, not a measurement.
- */
-const FOREIGN_TOKENIZER_SAFETY_FACTOR = 0.8
-
-/**
- * The token ceiling batching should enforce for a model.
+ * True when a model's tokens cannot be counted exactly.
  *
  * Batching measures with tiktoken, which only has encodings for OpenAI models —
  * every other id falls back to `cl100k_base`, so a Gemini, Cohere, or Mistral
- * limit would be enforced in OpenAI token units. Counting a foreign model's
- * text with tiktoken is an approximation, so its ceiling is discounted rather
- * than trusted exactly.
- *
- * The discount is deliberately one-sided: overshooting means the provider
- * rejects the whole request, while undershooting only trims a text that was
- * already at the limit.
+ * ceiling is enforced in approximate units. The ceiling is still applied
+ * exactly as declared: discounting it to absorb the error would truncate valid
+ * content silently, which is worse than the alternative it guards against. An
+ * undercount surfaces as a provider rejection, which is visible and
+ * actionable; silently shortening an embedding's input is not.
  */
-export function resolveBatchTokenCeiling(info: EmbeddingModelInfo): number {
-  if (info.tokenizerProvider === 'openai') return info.maxInputTokens
-  return Math.floor(info.maxInputTokens * FOREIGN_TOKENIZER_SAFETY_FACTOR)
+export function hasApproximateTokenCount(info: EmbeddingModelInfo): boolean {
+  return info.tokenizerProvider !== 'openai'
 }
 
 export function resolveDimensions(info: EmbeddingModelInfo, requested?: number): number {

apps/sim/lib/embeddings/client.ts

@@ -5,7 +5,7 @@ import {
   DEFAULT_EMBEDDING_MODEL,
   type EmbeddingModelInfo,
   getEmbeddingModelInfo,
-  resolveBatchTokenCeiling,
+  hasApproximateTokenCount,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 import { resolveProviderKey } from '@/lib/embeddings/keys'
@@ -193,17 +193,27 @@ export async function embed(texts: string[], options: EmbedOptions): Promise<Emb
   const modelInputs = options.projectInputs ? options.projectInputs(texts) : texts
 
   /**
-   * Batched against the selected model's own ceiling rather than one shared
-   * constant. A value that is too high sends oversized input the provider
-   * rejects; one that is too low silently drops content the provider would have
-   * accepted. Using the per-input ceiling as the per-batch budget also keeps
-   * every individual text within it.
+   * Batched against the selected model's own declared ceiling, exactly as
+   * declared. One shared constant sent oversized input to models with a lower
+   * limit and discarded content models with a higher one accept; discounting
+   * the ceiling to absorb tokenizer error would reintroduce the second harm.
+   *
+   * `batchByTokenLimit` truncates any text above the ceiling, and does so
+   * silently, so warn first: a shortened embedding input is otherwise
+   * indistinguishable from a good one, both to the caller and in the vector.
    */
-  const tokenBatches = batchByTokenLimit(
-    modelInputs,
-    resolveBatchTokenCeiling(provider.info),
-    model
-  )
+  const ceiling = provider.info.maxInputTokens
+  for (const text of modelInputs) {
+    if (estimateTokenCount(text, provider.info.tokenizerProvider).count <= ceiling) continue
+    logger.warn('Embedding input exceeds the model token limit and will be truncated', {
+      model,
+      maxInputTokens: ceiling,
+      chars: text.length,
+      approximateTokenCount: hasApproximateTokenCount(provider.info),
+    })
+  }
+
+  const tokenBatches = batchByTokenLimit(modelInputs, ceiling, model)
   const itemLimit = provider.adapter.maxItemsPerRequest ?? provider.info.maxItemsPerRequest
   const batches = itemLimit
     ? tokenBatches.flatMap((batch) => splitByItemLimit(batch, itemLimit))

apps/sim/lib/embeddings/index.ts

@@ -8,8 +8,8 @@ export {
   getEmbeddingModelInfo,
   getKbEligibleModels,
   getModelsForProvider,
+  hasApproximateTokenCount,
   KB_EMBEDDING_DIMENSIONS,
-  resolveBatchTokenCeiling,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 export { EmbeddingAPIError, embed } from '@/lib/embeddings/client'