fix(knowledge): bill only platform embedding tokens

Author: TheodoreSpeaks

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

@@ -286,6 +286,7 @@ describe('embed', () => {
     })
 
     expect(result.isBYOK).toBe(true)
+    expect(result.billableTokens).toBe(0)
   })
 
   it('uses OpenRouter as an explicit transport for an OpenAI catalog model', async () => {
@@ -455,6 +456,8 @@ describe('knowledge embedding transport fallback', () => {
     })
     expect(result).toMatchObject({
       embeddings: [[1, 2]],
+      billableTokens: 3,
+      isBYOK: false,
       modelName: 'text-embedding-3-small',
       dimensions: 1536,
     })
@@ -571,6 +574,7 @@ describe('knowledge embedding transport fallback', () => {
   it('falls back only the failed batch and retains successful provider work', async () => {
     vi.useFakeTimers()
     setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' })
+    mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-openai-test', isBYOK: true })
     const firstInput = `first ${'word '.repeat(5000)}`
     const secondInput = `second ${'word '.repeat(5000)}`
     fetchMock.mockImplementation(async (url, init) => {
@@ -582,7 +586,11 @@ describe('knowledge embedding transport fallback', () => {
       return jsonResponse(openAIBody([[input.startsWith('first') ? 1 : 2]], 3))
     })
 
-    const pending = embedKnowledgeForDeployment([firstInput, secondInput], options, false)
+    const pending = embedKnowledgeForDeployment(
+      [firstInput, secondInput],
+      { ...options, workspaceId: 'workspace-1' },
+      false
+    )
     await vi.runAllTimersAsync()
     const result = await pending
 
@@ -593,6 +601,8 @@ describe('knowledge embedding transport fallback', () => {
     expect(fetchMock).toHaveBeenCalledTimes(6)
     expect(result.embeddings).toEqual([[1], [2]])
     expect(result.totalTokens).toBe(6)
+    expect(result.billableTokens).toBe(3)
+    expect(result.isBYOK).toBe(false)
   })
 
   it('classifies only transient embedding failures for failover', () => {

apps/sim/lib/embeddings/client.ts

@@ -293,6 +293,7 @@ async function embedWithProvider(
   return {
     embeddings,
     totalTokens,
+    billableTokens: provider.isBYOK ? 0 : totalTokens,
     isBYOK: provider.isBYOK,
     modelName: provider.modelName,
     pricingId: provider.info.pricingId,
@@ -472,14 +473,16 @@ export async function embedKnowledgeForDeployment(
   const usedProviders = batchResults.map((batch) => batch.provider)
   const metadataProvider = usedProviders[0] ?? defaultProvider
   const modelNames = new Set(usedProviders.map((provider) => provider.modelName))
+  const billableTokens = batchResults.reduce(
+    (sum, batch) => sum + (batch.provider.isBYOK ? 0 : batch.totalTokens),
+    0
+  )
 
   return {
     embeddings,
     totalTokens,
-    isBYOK:
-      usedProviders.length > 0
-        ? usedProviders.every((provider) => provider.isBYOK)
-        : metadataProvider.isBYOK,
+    billableTokens,
+    isBYOK: usedProviders.length > 0 ? billableTokens === 0 : metadataProvider.isBYOK,
     modelName: modelNames.size > 1 ? model : metadataProvider.modelName,
     pricingId: info.pricingId,
     dimensions,

apps/sim/lib/embeddings/types.ts

@@ -108,7 +108,9 @@ export interface EmbedOptions {
 export interface EmbedResult {
   embeddings: number[][]
   totalTokens: number
-  /** True when a workspace-owned key was used, meaning Sim does not bill for it. */
+  /** Tokens processed with a Sim-funded key and therefore eligible for billing. */
+  billableTokens: number
+  /** True when every successful embedding used a caller- or workspace-owned key. */
   isBYOK: boolean
   /** Model name as sent to the provider. */
   modelName: string

apps/sim/lib/knowledge/documents/service.ts

@@ -924,8 +924,7 @@ export async function processDocumentAsync(
         .where(eq(document.id, documentId))
       return
     }
-    let totalEmbeddingTokens = 0
-    let embeddingIsBYOK = false
+    let billableEmbeddingTokens = 0
     let embeddingModelName = kbEmbeddingModel
     let embeddingPricingId = kbEmbeddingModel
 
@@ -992,17 +991,15 @@ export async function processDocumentAsync(
               logger.info(`[${documentId}] Processing embedding batch ${batchNum}/${totalBatches}`)
               const {
                 embeddings: batchEmbeddings,
-                totalTokens: batchTokens,
-                isBYOK,
+                billableTokens: batchBillableTokens,
                 modelName,
                 pricingId,
               } = await generateEmbeddings(batch, kbEmbeddingModel, ctx.workspaceId)
               for (const emb of batchEmbeddings) {
                 embeddings.push(emb)
               }
-              totalEmbeddingTokens += batchTokens
+              billableEmbeddingTokens += batchBillableTokens
               if (i === 0) {
-                embeddingIsBYOK = isBYOK
                 embeddingModelName = modelName
                 embeddingPricingId = pricingId
               }
@@ -1136,12 +1133,12 @@ export async function processDocumentAsync(
     const processingTime = Date.now() - startTime
     logger.info(`[${documentId}] Successfully processed document in ${processingTime}ms`)
 
-    if (!embeddingIsBYOK && totalEmbeddingTokens > 0) {
+    if (billableEmbeddingTokens > 0) {
       try {
         const costMultiplier = getCostMultiplier()
         const { total: cost } = calculateCost(
           embeddingPricingId,
-          totalEmbeddingTokens,
+          billableEmbeddingTokens,
           0,
           false,
           costMultiplier
@@ -1158,7 +1155,7 @@ export async function processDocumentAsync(
                 description: embeddingModelName,
                 cost,
                 sourceReference: `knowledge-document:${documentId}:${startTime}`,
-                metadata: { inputTokens: totalEmbeddingTokens, outputTokens: 0 },
+                metadata: { inputTokens: billableEmbeddingTokens, outputTokens: 0 },
               },
             ],
           })
@@ -1170,7 +1167,7 @@ export async function processDocumentAsync(
         } else {
           logger.warn(
             `[${documentId}] Embedding model "${embeddingModelName}" has no pricing entry — billing skipped`,
-            { totalEmbeddingTokens, embeddingModelName }
+            { billableEmbeddingTokens, embeddingModelName }
           )
         }
       } catch (billingError) {

apps/sim/lib/knowledge/embeddings.ts

@@ -46,6 +46,7 @@ export function getConfiguredEmbeddingModel(): string {
 export interface GenerateEmbeddingsResult {
   embeddings: number[][]
   totalTokens: number
+  billableTokens: number
   isBYOK: boolean
   modelName: string
   /** Pricing identifier for use with calculateCost / EMBEDDING_MODEL_PRICING. */
@@ -76,6 +77,7 @@ export async function generateEmbeddings(
   return {
     embeddings: result.embeddings,
     totalTokens: result.totalTokens,
+    billableTokens: result.billableTokens,
     isBYOK: result.isBYOK,
     modelName: result.modelName,
     pricingId: result.pricingId,