chore: import upstream snapshot with attribution

apps/sim/lib/knowledge/embeddings.ts

@@ -0,0 +1,451 @@
+import { createLogger } from '@sim/logger'
+import { getErrorMessage } from '@sim/utils/errors'
+import { getBYOKKey } from '@/lib/api-key/byok'
+import { recordUsage } from '@/lib/billing/core/usage-log'
+import { getRotatingApiKey } from '@/lib/core/config/api-keys'
+import { env, envNumber } from '@/lib/core/config/env'
+import { isRetryableError, retryWithExponentialBackoff } from '@/lib/knowledge/documents/utils'
+import {
+  DEFAULT_EMBEDDING_MODEL,
+  EMBEDDING_DIMENSIONS,
+  getEmbeddingModelInfo,
+  SUPPORTED_EMBEDDING_MODELS,
+  type TokenizerProviderId,
+} from '@/lib/knowledge/embedding-models'
+import { batchByTokenLimit, estimateTokenCount } from '@/lib/tokenization'
+import { calculateCost } from '@/providers/utils'
+
+const logger = createLogger('EmbeddingUtils')
+
+const MAX_TOKENS_PER_REQUEST = 8000
+const MAX_CONCURRENT_BATCHES = envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 50)
+const EMBEDDING_REQUEST_TIMEOUT_MS = 60_000
+
+export { EMBEDDING_DIMENSIONS } from '@/lib/knowledge/embedding-models'
+
+class EmbeddingAPIError extends Error {
+  public status: number
+
+  constructor(message: string, status: number) {
+    super(message)
+    this.name = 'EmbeddingAPIError'
+    this.status = status
+  }
+}
+
+export type EmbeddingInputType = 'document' | 'query'
+
+interface ProviderRequest {
+  apiUrl: string
+  headers: Record<string, string>
+  body: unknown
+  parse: (json: unknown) => number[][]
+}
+
+interface ResolvedProvider {
+  modelName: string
+  pricingId: string
+  isBYOK: boolean
+  /** Tokenizer used to estimate tokens when the API does not return a usage field. */
+  tokenizerProvider: TokenizerProviderId
+  /** Hard per-request item cap enforced by the provider (e.g. Gemini caps at 100). */
+  maxItemsPerRequest?: number
+  buildRequest: (inputs: string[], inputType: EmbeddingInputType) => ProviderRequest
+}
+
+/** Gemini's `batchEmbedContents` rejects requests with more than 100 items. */
+const GEMINI_MAX_ITEMS_PER_REQUEST = 100
+
+async function resolveOpenAIKey(workspaceId?: string | null): Promise<{
+  apiKey: string
+  isBYOK: boolean
+}> {
+  if (workspaceId) {
+    const byokResult = await getBYOKKey(workspaceId, 'openai')
+    if (byokResult) {
+      logger.info('Using workspace BYOK key for OpenAI embeddings')
+      return { apiKey: byokResult.apiKey, isBYOK: true }
+    }
+  }
+  if (env.OPENAI_API_KEY) {
+    return { apiKey: env.OPENAI_API_KEY, isBYOK: false }
+  }
+  try {
+    return { apiKey: getRotatingApiKey('openai'), isBYOK: false }
+  } catch {
+    throw new Error('OPENAI_API_KEY is not configured')
+  }
+}
+
+async function resolveGeminiKey(workspaceId?: string | null): Promise<{
+  apiKey: string
+  isBYOK: boolean
+}> {
+  if (workspaceId) {
+    const byokResult = await getBYOKKey(workspaceId, 'google')
+    if (byokResult) {
+      logger.info('Using workspace BYOK key for Gemini embeddings')
+      return { apiKey: byokResult.apiKey, isBYOK: true }
+    }
+  }
+  if (env.GEMINI_API_KEY) {
+    return { apiKey: env.GEMINI_API_KEY, isBYOK: false }
+  }
+  try {
+    return { apiKey: getRotatingApiKey('gemini'), isBYOK: false }
+  } catch {
+    throw new Error(
+      'GEMINI_API_KEY (or GEMINI_API_KEY_1/2/3 for rotation) must be configured for Gemini embeddings'
+    )
+  }
+}
+
+function buildOpenAIProvider(modelName: string, apiKey: string): ResolvedProvider['buildRequest'] {
+  return (inputs) => ({
+    apiUrl: 'https://api.openai.com/v1/embeddings',
+    headers: {
+      Authorization: `Bearer ${apiKey}`,
+      'Content-Type': 'application/json',
+    },
+    body: {
+      input: inputs,
+      model: modelName,
+      encoding_format: 'float',
+      dimensions: EMBEDDING_DIMENSIONS,
+    },
+    parse: (json) => {
+      const data = json as { data: Array<{ embedding: number[] }> }
+      return data.data.map((item) => item.embedding)
+    },
+  })
+}
+
+function buildAzureOpenAIProvider(
+  deployment: string,
+  apiKey: string,
+  endpoint: string,
+  apiVersion: string
+): ResolvedProvider['buildRequest'] {
+  return (inputs) => ({
+    apiUrl: `${endpoint}/openai/deployments/${deployment}/embeddings?api-version=${apiVersion}`,
+    headers: {
+      'api-key': apiKey,
+      'Content-Type': 'application/json',
+    },
+    body: {
+      input: inputs,
+      encoding_format: 'float',
+      dimensions: EMBEDDING_DIMENSIONS,
+    },
+    parse: (json) => {
+      const data = json as { data: Array<{ embedding: number[] }> }
+      return data.data.map((item) => item.embedding)
+    },
+  })
+}
+
+/**
+ * Gemini does NOT auto-normalize embeddings when `outputDimensionality` is set below the
+ * native 3072 dimension on `gemini-embedding-001`. Manually L2-normalize so cosine and
+ * inner-product similarity work correctly.
+ */
+function l2Normalize(vector: number[]): number[] {
+  let sumSquares = 0
+  for (const v of vector) sumSquares += v * v
+  const norm = Math.sqrt(sumSquares)
+  if (norm === 0) return vector
+  return vector.map((v) => v / norm)
+}
+
+function buildGeminiProvider(modelName: string, apiKey: string): ResolvedProvider['buildRequest'] {
+  return (inputs, inputType) => ({
+    apiUrl: `https://generativelanguage.googleapis.com/v1beta/models/${modelName}:batchEmbedContents`,
+    headers: {
+      'Content-Type': 'application/json',
+      'x-goog-api-key': apiKey,
+    },
+    body: {
+      requests: inputs.map((text) => ({
+        model: `models/${modelName}`,
+        content: { parts: [{ text }] },
+        taskType: inputType === 'query' ? 'RETRIEVAL_QUERY' : 'RETRIEVAL_DOCUMENT',
+        outputDimensionality: EMBEDDING_DIMENSIONS,
+      })),
+    },
+    parse: (json) => {
+      const data = json as { embeddings: Array<{ values: number[] }> }
+      return data.embeddings.map((item) => l2Normalize(item.values))
+    },
+  })
+}
+
+/**
+ * Returns the embedding model to use for new knowledge bases.
+ * Sourced from the `KB_EMBEDDING_MODEL` env var; falls back to the default if
+ * unset or set to an unsupported model.
+ */
+export function getConfiguredEmbeddingModel(): string {
+  const configured = env.KB_EMBEDDING_MODEL
+  if (configured && SUPPORTED_EMBEDDING_MODELS[configured]) {
+    return configured
+  }
+  if (configured) {
+    logger.warn(
+      `KB_EMBEDDING_MODEL="${configured}" is not a supported embedding model — falling back to ${DEFAULT_EMBEDDING_MODEL}`
+    )
+  }
+  return DEFAULT_EMBEDDING_MODEL
+}
+
+async function resolveProvider(
+  embeddingModel: string,
+  workspaceId?: string | null
+): Promise<ResolvedProvider> {
+  const azureApiKey = env.AZURE_OPENAI_API_KEY
+  const azureEndpoint = env.AZURE_OPENAI_ENDPOINT
+  const azureApiVersion = env.AZURE_OPENAI_API_VERSION
+  const isOpenAIModel = SUPPORTED_EMBEDDING_MODELS[embeddingModel]?.provider === 'openai'
+  /**
+   * Azure deployment names default to the embedding model name when
+   * `KB_OPENAI_MODEL_NAME` is unset — this matches the pre-existing
+   * convention where deployments are named after the model they host.
+   */
+  const azureDeploymentName = env.KB_OPENAI_MODEL_NAME || embeddingModel
+  const useAzure = Boolean(isOpenAIModel && azureApiKey && azureEndpoint && azureApiVersion)
+
+  const info = getEmbeddingModelInfo(embeddingModel)
+
+  if (useAzure) {
+    return {
+      modelName: azureDeploymentName,
+      pricingId: info.pricingId,
+      isBYOK: false,
+      tokenizerProvider: info.tokenizerProvider,
+      buildRequest: buildAzureOpenAIProvider(
+        azureDeploymentName,
+        azureApiKey!,
+        azureEndpoint!,
+        azureApiVersion!
+      ),
+    }
+  }
+
+  if (info.provider === 'openai') {
+    const { apiKey, isBYOK } = await resolveOpenAIKey(workspaceId)
+    return {
+      modelName: embeddingModel,
+      pricingId: info.pricingId,
+      isBYOK,
+      tokenizerProvider: info.tokenizerProvider,
+      buildRequest: buildOpenAIProvider(embeddingModel, apiKey),
+    }
+  }
+
+  if (info.provider === 'gemini') {
+    const { apiKey, isBYOK } = await resolveGeminiKey(workspaceId)
+    return {
+      modelName: embeddingModel,
+      pricingId: info.pricingId,
+      isBYOK,
+      tokenizerProvider: info.tokenizerProvider,
+      maxItemsPerRequest: GEMINI_MAX_ITEMS_PER_REQUEST,
+      buildRequest: buildGeminiProvider(embeddingModel, apiKey),
+    }
+  }
+
+  throw new Error(`Unknown embedding provider for model ${embeddingModel}`)
+}
+
+async function callEmbeddingAPI(
+  inputs: string[],
+  provider: ResolvedProvider,
+  inputType: EmbeddingInputType
+): Promise<{ embeddings: number[][]; totalTokens: number }> {
+  return retryWithExponentialBackoff(
+    async () => {
+      const request = provider.buildRequest(inputs, inputType)
+
+      const controller = new AbortController()
+      const timeout = setTimeout(() => controller.abort(), EMBEDDING_REQUEST_TIMEOUT_MS)
+
+      const response = await fetch(request.apiUrl, {
+        method: 'POST',
+        headers: request.headers,
+        body: JSON.stringify(request.body),
+        signal: controller.signal,
+      }).finally(() => clearTimeout(timeout))
+
+      if (!response.ok) {
+        const errorText = await response.text()
+        throw new EmbeddingAPIError(
+          `Embedding API failed: ${response.status} ${response.statusText} - ${errorText}`,
+          response.status
+        )
+      }
+
+      const json = await response.json()
+      const embeddings = request.parse(json)
+      const usage = (json as { usage?: { total_tokens?: number } }).usage
+      const totalTokens =
+        usage?.total_tokens ??
+        // Gemini does not return usage.total_tokens — estimate with the provider's tokenizer
+        inputs.reduce(
+          (sum, text) => sum + estimateTokenCount(text, provider.tokenizerProvider).count,
+          0
+        )
+
+      return { embeddings, totalTokens }
+    },
+    {
+      maxRetries: 3,
+      initialDelayMs: 1000,
+      maxDelayMs: 10000,
+      retryCondition: (error: unknown) => {
+        if (error instanceof EmbeddingAPIError) {
+          return error.status === 429 || error.status >= 500
+        }
+        return isRetryableError(error)
+      },
+    }
+  )
+}
+
+function splitByItemLimit<T>(items: T[], limit: number): T[][] {
+  if (items.length <= limit) return [items]
+  const result: T[][] = []
+  for (let i = 0; i < items.length; i += limit) {
+    result.push(items.slice(i, i + limit))
+  }
+  return result
+}
+
+async function processWithConcurrency<T, R>(
+  items: T[],
+  concurrency: number,
+  processor: (item: T, index: number) => Promise<R>
+): Promise<R[]> {
+  const results: R[] = new Array(items.length)
+  let currentIndex = 0
+
+  const workers = Array.from({ length: Math.min(concurrency, items.length) }, async () => {
+    while (currentIndex < items.length) {
+      const index = currentIndex++
+      results[index] = await processor(items[index], index)
+    }
+  })
+
+  await Promise.all(workers)
+  return results
+}
+
+export interface GenerateEmbeddingsResult {
+  embeddings: number[][]
+  totalTokens: number
+  isBYOK: boolean
+  modelName: string
+  /** Pricing identifier for use with calculateCost / EMBEDDING_MODEL_PRICING. */
+  pricingId: string
+}
+
+/**
+ * Generate embeddings for multiple texts with token-aware batching and parallel processing.
+ */
+export async function generateEmbeddings(
+  texts: string[],
+  embeddingModel: string = DEFAULT_EMBEDDING_MODEL,
+  workspaceId?: string | null
+): Promise<GenerateEmbeddingsResult> {
+  const provider = await resolveProvider(embeddingModel, workspaceId)
+
+  const tokenBatches = batchByTokenLimit(texts, MAX_TOKENS_PER_REQUEST, embeddingModel)
+  const batches = provider.maxItemsPerRequest
+    ? tokenBatches.flatMap((batch) => splitByItemLimit(batch, provider.maxItemsPerRequest!))
+    : tokenBatches
+
+  const batchResults = await processWithConcurrency(
+    batches,
+    MAX_CONCURRENT_BATCHES,
+    async (batch, i) => {
+      try {
+        return await callEmbeddingAPI(batch, provider, 'document')
+      } catch (error) {
+        logger.error(`Failed to generate embeddings for batch ${i + 1}/${batches.length}:`, error)
+        throw error
+      }
+    }
+  )
+
+  const allEmbeddings: number[][] = []
+  let totalTokens = 0
+  for (const batch of batchResults) {
+    for (const emb of batch.embeddings) {
+      allEmbeddings.push(emb)
+    }
+    totalTokens += batch.totalTokens
+  }
+
+  return {
+    embeddings: allEmbeddings,
+    totalTokens,
+    isBYOK: provider.isBYOK,
+    modelName: provider.modelName,
+    pricingId: provider.pricingId,
+  }
+}
+
+/**
+ * Generate embedding for a single search query.
+ */
+export async function generateSearchEmbedding(
+  query: string,
+  embeddingModel: string = DEFAULT_EMBEDDING_MODEL,
+  workspaceId?: string | null
+): Promise<{ embedding: number[]; isBYOK: boolean }> {
+  const provider = await resolveProvider(embeddingModel, workspaceId)
+
+  logger.info(`Using ${provider.modelName} for search embedding generation`)
+
+  const { embeddings } = await callEmbeddingAPI([query], provider, 'query')
+  return { embedding: embeddings[0], isBYOK: provider.isBYOK }
+}
+
+/**
+ * Records a query embedding's hosted-key cost for callers that generate a search
+ * embedding directly, outside the metered `/api/knowledge/search` route (e.g. the
+ * v1 search API and copilot KB search). No-ops for BYOK (no Sim cost) or when
+ * there is no workspace to attribute to. Best-effort: never throws.
+ */
+export async function recordSearchEmbeddingUsage(params: {
+  userId: string
+  workspaceId?: string | null
+  embeddingModel: string
+  query: string
+  isBYOK: boolean
+  sourceReference: string
+}): Promise<void> {
+  const { userId, workspaceId, embeddingModel, query, isBYOK, sourceReference } = params
+  if (isBYOK || !workspaceId) return
+  try {
+    const { count } = estimateTokenCount(
+      query,
+      getEmbeddingModelInfo(embeddingModel).tokenizerProvider
+    )
+    const cost = calculateCost(embeddingModel, count, 0, false)
+    if (!cost || cost.total <= 0) return
+    await recordUsage({
+      userId,
+      workspaceId,
+      entries: [
+        {
+          category: 'model',
+          source: 'knowledge-base',
+          description: embeddingModel,
+          cost: cost.total,
+          sourceReference,
+        },
+      ],
+    })
+  } catch (error) {
+    logger.warn('Failed to record search embedding usage', { error: getErrorMessage(error) })
+  }
+}