feat(byok): byok for hosted model capabilities

apps/sim/app/api/knowledge/[id]/documents/[documentId]/chunks/[chunkId]/route.ts

@@ -100,7 +100,12 @@ export async function PUT(
     try {
       const validatedData = UpdateChunkSchema.parse(body)
 
-      const updatedChunk = await updateChunk(chunkId, validatedData, requestId)
+      const updatedChunk = await updateChunk(
+        chunkId,
+        validatedData,
+        requestId,
+        accessCheck.knowledgeBase?.workspaceId
+      )
 
       logger.info(
         `[${requestId}] Chunk updated: ${chunkId} in document ${documentId} in knowledge base ${knowledgeBaseId}`

apps/sim/app/api/knowledge/[id]/documents/[documentId]/chunks/route.ts

@@ -184,7 +184,8 @@ export async function POST(
         documentId,
         docTags,
         validatedData,
-        requestId
+        requestId,
+        accessCheck.knowledgeBase?.workspaceId
       )
 
       let cost = null

apps/sim/app/api/knowledge/search/route.ts

@@ -183,11 +183,11 @@ export async function POST(request: NextRequest) {
         )
       }
 
-      // Generate query embedding only if query is provided
+      const workspaceId = accessChecks.find((ac) => ac?.hasAccess)?.knowledgeBase?.workspaceId
+
       const hasQuery = validatedData.query && validatedData.query.trim().length > 0
-      // Start embedding generation early and await when needed
       const queryEmbeddingPromise = hasQuery
-        ? generateSearchEmbedding(validatedData.query!)
+        ? generateSearchEmbedding(validatedData.query!, undefined, workspaceId)
         : Promise.resolve(null)
 
       // Check if any requested knowledge bases were not accessible

apps/sim/app/api/knowledge/utils.ts

@@ -113,7 +113,7 @@ export type KnowledgeBaseAccessCheck = KnowledgeBaseAccessResult | KnowledgeBase
 export interface DocumentAccessResult {
   hasAccess: true
   document: DocumentData
-  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId'>
+  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId' | 'workspaceId'>
 }
 
 export interface DocumentAccessDenied {
@@ -128,7 +128,7 @@ export interface ChunkAccessResult {
   hasAccess: true
   chunk: EmbeddingData
   document: DocumentData
-  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId'>
+  knowledgeBase: Pick<KnowledgeBaseData, 'id' | 'userId' | 'workspaceId'>
 }
 
 export interface ChunkAccessDenied {

apps/sim/app/api/providers/route.ts

@@ -2,12 +2,12 @@ import { db } from '@sim/db'
 import { account } from '@sim/db/schema'
 import { eq } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
+import { getApiKeyWithBYOK } from '@/lib/api-key/byok'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { createLogger } from '@/lib/logs/console/logger'
 import { refreshTokenIfNeeded } from '@/app/api/auth/oauth/utils'
 import type { StreamingExecution } from '@/executor/types'
 import { executeProviderRequest } from '@/providers'
-import { getApiKey } from '@/providers/utils'
 
 const logger = createLogger('ProvidersAPI')
 
@@ -81,11 +81,14 @@ export async function POST(request: NextRequest) {
     })
 
     let finalApiKey: string
+    let isBYOK = body.isBYOK ?? false
     try {
       if (provider === 'vertex' && vertexCredential) {
         finalApiKey = await resolveVertexCredential(requestId, vertexCredential)
       } else {
-        finalApiKey = getApiKey(provider, model, apiKey)
+        const result = await getApiKeyWithBYOK(provider, model, workspaceId, apiKey)
+        finalApiKey = result.apiKey
+        isBYOK = result.isBYOK
       }
     } catch (error) {
       logger.error(`[${requestId}] Failed to get API key:`, {
@@ -106,9 +109,9 @@ export async function POST(request: NextRequest) {
       model,
       workflowId,
       hasApiKey: !!finalApiKey,
+      isBYOK,
     })
 
-    // Execute provider request directly with the managed key
     const response = await executeProviderRequest(provider, {
       model,
       systemPrompt,
@@ -117,6 +120,7 @@ export async function POST(request: NextRequest) {
       temperature,
       maxTokens,
       apiKey: finalApiKey,
+      isBYOK,
       azureEndpoint,
       azureApiVersion,
       vertexProject,

apps/sim/app/api/tools/search/route.ts

@@ -1,5 +1,6 @@
 import { type NextRequest, NextResponse } from 'next/server'
 import { z } from 'zod'
+import { getBYOKKey } from '@/lib/api-key/byok'
 import { checkHybridAuth } from '@/lib/auth/hybrid'
 import { SEARCH_TOOL_COST } from '@/lib/billing/constants'
 import { env } from '@/lib/core/config/env'
@@ -10,6 +11,7 @@ const logger = createLogger('search')
 
 const SearchRequestSchema = z.object({
   query: z.string().min(1),
+  workspaceId: z.string().optional(),
 })
 
 export const maxDuration = 60
@@ -39,8 +41,20 @@ export async function POST(request: NextRequest) {
     const body = await request.json()
     const validated = SearchRequestSchema.parse(body)
 
-    if (!env.EXA_API_KEY) {
-      logger.error(`[${requestId}] EXA_API_KEY not configured`)
+    let exaApiKey = env.EXA_API_KEY
+    let isBYOK = false
+
+    if (validated.workspaceId) {
+      const byokResult = await getBYOKKey(validated.workspaceId, 'exa')
+      if (byokResult) {
+        exaApiKey = byokResult.apiKey
+        isBYOK = true
+        logger.info(`[${requestId}] Using workspace BYOK key for Exa search`)
+      }
+    }
+
+    if (!exaApiKey) {
+      logger.error(`[${requestId}] No Exa API key available`)
       return NextResponse.json(
         { success: false, error: 'Search service not configured' },
         { status: 503 }
@@ -50,14 +64,15 @@ export async function POST(request: NextRequest) {
     logger.info(`[${requestId}] Executing search`, {
       userId,
       query: validated.query,
+      isBYOK,
     })
 
     const result = await executeTool('exa_search', {
       query: validated.query,
       type: 'auto',
       useAutoprompt: true,
       highlights: true,
-      apiKey: env.EXA_API_KEY,
+      apiKey: exaApiKey,
     })
 
     if (!result.success) {
@@ -85,7 +100,7 @@ export async function POST(request: NextRequest) {
     const cost = {
       input: 0,
       output: 0,
-      total: SEARCH_TOOL_COST,
+      total: isBYOK ? 0 : SEARCH_TOOL_COST,
       tokens: {
         input: 0,
         output: 0,
@@ -104,6 +119,7 @@ export async function POST(request: NextRequest) {
       userId,
       resultCount: results.length,
       cost: cost.total,
+      isBYOK,
     })
 
     return NextResponse.json({

apps/sim/app/api/wand/route.ts

@@ -3,6 +3,7 @@ import { userStats, workflow } from '@sim/db/schema'
 import { eq, sql } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import OpenAI, { AzureOpenAI } from 'openai'
+import { getBYOKKey } from '@/lib/api-key/byok'
 import { getSession } from '@/lib/auth'
 import { logModelUsage } from '@/lib/billing/core/usage-log'
 import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
@@ -75,7 +76,8 @@ async function updateUserStatsForWand(
     completion_tokens?: number
     total_tokens?: number
   },
-  requestId: string
+  requestId: string,
+  isBYOK = false
 ): Promise<void> {
   if (!isBillingEnabled) {
     logger.debug(`[${requestId}] Billing is disabled, skipping wand usage cost update`)
@@ -93,21 +95,24 @@ async function updateUserStatsForWand(
     const completionTokens = usage.completion_tokens || 0
 
     const modelName = useWandAzure ? wandModelName : 'gpt-4o'
-    const pricing = getModelPricing(modelName)
-
-    const costMultiplier = getCostMultiplier()
-    let modelCost = 0
+    let costToStore = 0
+
+    if (!isBYOK) {
+      const pricing = getModelPricing(modelName)
+      const costMultiplier = getCostMultiplier()
+      let modelCost = 0
+
+      if (pricing) {
+        const inputCost = (promptTokens / 1000000) * pricing.input
+        const outputCost = (completionTokens / 1000000) * pricing.output
+        modelCost = inputCost + outputCost
+      } else {
+        modelCost = (promptTokens / 1000000) * 0.005 + (completionTokens / 1000000) * 0.015
+      }
 
-    if (pricing) {
-      const inputCost = (promptTokens / 1000000) * pricing.input
-      const outputCost = (completionTokens / 1000000) * pricing.output
-      modelCost = inputCost + outputCost
-    } else {
-      modelCost = (promptTokens / 1000000) * 0.005 + (completionTokens / 1000000) * 0.015
+      costToStore = modelCost * costMultiplier
     }
 
-    const costToStore = modelCost * costMultiplier
-
     await db
       .update(userStats)
       .set({
@@ -122,6 +127,7 @@ async function updateUserStatsForWand(
       userId,
       tokensUsed: totalTokens,
       costAdded: costToStore,
+      isBYOK,
     })
 
     await logModelUsage({
@@ -149,14 +155,6 @@ export async function POST(req: NextRequest) {
     return NextResponse.json({ success: false, error: 'Unauthorized' }, { status: 401 })
   }
 
-  if (!client) {
-    logger.error(`[${requestId}] AI client not initialized. Missing API key.`)
-    return NextResponse.json(
-      { success: false, error: 'Wand generation service is not configured.' },
-      { status: 503 }
-    )
-  }
-
   try {
     const body = (await req.json()) as RequestBody
 
@@ -170,6 +168,7 @@ export async function POST(req: NextRequest) {
       )
     }
 
+    let workspaceId: string | null = null
     if (workflowId) {
       const [workflowRecord] = await db
         .select({ workspaceId: workflow.workspaceId, userId: workflow.userId })
@@ -182,6 +181,8 @@ export async function POST(req: NextRequest) {
         return NextResponse.json({ success: false, error: 'Workflow not found' }, { status: 404 })
       }
 
+      workspaceId = workflowRecord.workspaceId
+
       if (workflowRecord.workspaceId) {
         const permission = await verifyWorkspaceMembership(
           session.user.id,
@@ -199,6 +200,28 @@ export async function POST(req: NextRequest) {
       }
     }
 
+    let isBYOK = false
+    let activeClient = client
+    let byokApiKey: string | null = null
+
+    if (workspaceId && !useWandAzure) {
+      const byokResult = await getBYOKKey(workspaceId, 'openai')
+      if (byokResult) {
+        isBYOK = true
+        byokApiKey = byokResult.apiKey
+        activeClient = new OpenAI({ apiKey: byokResult.apiKey })
+        logger.info(`[${requestId}] Using BYOK OpenAI key for wand generation`)
+      }
+    }
+
+    if (!activeClient) {
+      logger.error(`[${requestId}] AI client not initialized. Missing API key.`)
+      return NextResponse.json(
+        { success: false, error: 'Wand generation service is not configured.' },
+        { status: 503 }
+      )
+    }
+
     const finalSystemPrompt =
       systemPrompt ||
       'You are a helpful AI assistant. Generate content exactly as requested by the user.'
@@ -241,7 +264,7 @@ export async function POST(req: NextRequest) {
         if (useWandAzure) {
           headers['api-key'] = azureApiKey!
         } else {
-          headers.Authorization = `Bearer ${openaiApiKey}`
+          headers.Authorization = `Bearer ${byokApiKey || openaiApiKey}`
         }
 
         logger.debug(`[${requestId}] Making streaming request to: ${apiUrl}`)
@@ -310,7 +333,7 @@ export async function POST(req: NextRequest) {
                       logger.info(`[${requestId}] Received [DONE] signal`)
 
                       if (finalUsage) {
-                        await updateUserStatsForWand(session.user.id, finalUsage, requestId)
+                        await updateUserStatsForWand(session.user.id, finalUsage, requestId, isBYOK)
                       }
 
                       controller.enqueue(
@@ -395,7 +418,7 @@ export async function POST(req: NextRequest) {
       }
     }
 
-    const completion = await client.chat.completions.create({
+    const completion = await activeClient.chat.completions.create({
       model: useWandAzure ? wandModelName : 'gpt-4o',
       messages: messages,
       temperature: 0.3,
@@ -417,7 +440,7 @@ export async function POST(req: NextRequest) {
     logger.info(`[${requestId}] Wand generation successful`)
 
     if (completion.usage) {
-      await updateUserStatsForWand(session.user.id, completion.usage, requestId)
+      await updateUserStatsForWand(session.user.id, completion.usage, requestId, isBYOK)
     }
 
     return NextResponse.json({ success: true, content: generatedContent })