fix(embeddings): discount the batch ceiling when the tokenizer is foreign

Review round 4.

Batching measures with tiktoken, which only has encodings for OpenAI models —
every other id falls back to cl100k_base. Gemini's 2048, Cohere's 128k, and
Mistral's 8192 were therefore enforced in OpenAI token units, so an input near
one of those ceilings could still be rejected upstream or trimmed more than
needed.

A true fix needs per-provider tokenizers, which 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. So the ceiling is discounted
for foreign tokenizers rather than trusted exactly.

The discount is one-sided on purpose. Overshooting means the provider rejects
the whole request; undershooting only trims a text that was already at the
limit, so the margin errs toward the second.

resolveBatchTokenCeiling is a pure function tested directly, rather than
inferred from truncation behavior, so the guarantee holds per model as the
catalog grows.

Author: mzxchandra

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

@@ -9,6 +9,7 @@ import {
   getKbEligibleModels,
   getModelsForProvider,
   KB_EMBEDDING_DIMENSIONS,
+  resolveBatchTokenCeiling,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 import { EMBEDDING_MODEL_PRICING } from '@/providers/models'
@@ -100,6 +101,38 @@ 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)

apps/sim/lib/embeddings/catalog.ts

@@ -157,6 +157,31 @@ export function getKbEligibleModels(): string[] {
  * Resolves the dimensionality a request will actually produce, given an
  * 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.
+ *
+ * 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.
+ */
+export function resolveBatchTokenCeiling(info: EmbeddingModelInfo): number {
+  if (info.tokenizerProvider === 'openai') return info.maxInputTokens
+  return Math.floor(info.maxInputTokens * FOREIGN_TOKENIZER_SAFETY_FACTOR)
+}
+
 export function resolveDimensions(info: EmbeddingModelInfo, requested?: number): number {
   if (requested === undefined) return info.nativeDimensions
   if (!info.supportedDimensions?.includes(requested)) {

apps/sim/lib/embeddings/client.ts

@@ -5,6 +5,7 @@ import {
   DEFAULT_EMBEDDING_MODEL,
   type EmbeddingModelInfo,
   getEmbeddingModelInfo,
+  resolveBatchTokenCeiling,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 import { resolveProviderKey } from '@/lib/embeddings/keys'
@@ -198,7 +199,11 @@ export async function embed(texts: string[], options: EmbedOptions): Promise<Emb
    * accepted. Using the per-input ceiling as the per-batch budget also keeps
    * every individual text within it.
    */
-  const tokenBatches = batchByTokenLimit(modelInputs, provider.info.maxInputTokens, model)
+  const tokenBatches = batchByTokenLimit(
+    modelInputs,
+    resolveBatchTokenCeiling(provider.info),
+    model
+  )
   const itemLimit = provider.adapter.maxItemsPerRequest ?? provider.info.maxItemsPerRequest
   const batches = itemLimit
     ? tokenBatches.flatMap((batch) => splitByItemLimit(batch, itemLimit))

apps/sim/lib/embeddings/index.ts

@@ -9,6 +9,7 @@ export {
   getKbEligibleModels,
   getModelsForProvider,
   KB_EMBEDDING_DIMENSIONS,
+  resolveBatchTokenCeiling,
   resolveDimensions,
 } from '@/lib/embeddings/catalog'
 export { EmbeddingAPIError, embed } from '@/lib/embeddings/client'