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

import { createLogger } from '@sim/logger'
import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow'
import { getErrorMessage } from '@sim/utils/errors'
import { type NextRequest, NextResponse } from 'next/server'
import { knowledgeSearchBodySchema } from '@/lib/api/contracts/knowledge'
import { parseJsonBody, validationErrorResponse } from '@/lib/api/server'
import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
import { checkActorUsageLimits } from '@/lib/billing/calculations/usage-monitor'
import {
  checkAttributedUsageLimits,
  requireBillingAttributionHeader,
  resolveBillingAttribution,
  toBillingContext,
} from '@/lib/billing/core/billing-attribution'
import {
  checkAndBillOverageThreshold,
  checkAndBillPayerOverageThreshold,
} from '@/lib/billing/threshold-billing'
import { PlatformEvents } from '@/lib/core/telemetry'
import { generateRequestId } from '@/lib/core/utils/request'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { importDurableSecretProvenance } from '@/lib/execution/durable-secret-provenance'
import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants'
import { getEmbeddingModelInfo } from '@/lib/knowledge/embedding-models'
import {
  prepareKnowledgeModelInputProvenance,
  runWithKnowledgeModelInputProvenance,
} from '@/lib/knowledge/model-input-provenance'
import { rerank } from '@/lib/knowledge/reranker'
import { importKnowledgeSearchResultSecretProvenance } from '@/lib/knowledge/secret-provenance'
import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service'
import { buildUndefinedTagsError, validateTagValue } from '@/lib/knowledge/tags/utils'
import type { StructuredFilter } from '@/lib/knowledge/types'
import { estimateTokenCount } from '@/lib/tokenization/estimators'
import {
  executeKnowledgeSearch,
  generateSearchEmbedding,
  type SearchResult,
} from '@/app/api/knowledge/search/utils'
import { createKnowledgeRegistryResponse } from '@/app/api/knowledge/secret-provenance'
import { checkKnowledgeBaseAccess, type KnowledgeBaseAccessResult } from '@/app/api/knowledge/utils'
import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
import { getRerankModelPricing } from '@/providers/models'
import { calculateCost } from '@/providers/utils'

const logger = createLogger('VectorSearchAPI')

export const POST = withRouteHandler(async (request: NextRequest) => {
  const requestId = generateRequestId()

  try {
    const parsedBody = await parseJsonBody(request)
    if (!parsedBody.success) return parsedBody.response
    const body = parsedBody.data as Record<string, unknown>
    const { workflowId, skipUsageBilling, ...searchParams } = body

    const auth = await checkSessionOrInternalAuth(request, { requireWorkflowId: false })
    if (!auth.success || !auth.userId) {
      return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
    }
    const userId = auth.userId

    // Only the internal workflow tool may suppress route metering (it rolls the
    // cost into the executor's usage instead). Session/API-key callers cannot set
    // skipUsageBilling to dodge their own embedding/reranker charge.
    const shouldMeter = !(skipUsageBilling === true && auth.authType === AuthType.INTERNAL_JWT)

    if (workflowId) {
      const authorization = await authorizeWorkflowByWorkspacePermission({
        workflowId: workflowId as string,
        userId,
        action: 'read',
      })
      if (!authorization.allowed) {
        return NextResponse.json(
          { error: authorization.message || 'Access denied' },
          { status: authorization.status }
        )
      }
    }

    const validation = knowledgeSearchBodySchema.safeParse(searchParams)
    if (!validation.success) return validationErrorResponse(validation.error)
    const validatedData = validation.data

    const knowledgeBaseIds = Array.isArray(validatedData.knowledgeBaseIds)
      ? validatedData.knowledgeBaseIds
      : [validatedData.knowledgeBaseIds]

    const accessChecks = await Promise.all(
      knowledgeBaseIds.map((kbId) => checkKnowledgeBaseAccess(kbId, userId))
    )
    const accessibleKbIds: string[] = knowledgeBaseIds.filter(
      (_, idx) => accessChecks[idx]?.hasAccess
    )

    let structuredFilters: StructuredFilter[] = []

    if (validatedData.tagFilters && accessibleKbIds.length > 0) {
      const kbTagDefs = await Promise.all(
        accessibleKbIds.map(async (kbId) => ({
          kbId,
          tagDefs: await getDocumentTagDefinitions(kbId),
        }))
      )

      const displayNameToTagDef: Record<string, { tagSlot: string; fieldType: string }> = {}
      for (const { kbId, tagDefs } of kbTagDefs) {
        const perKbMap = new Map(
          tagDefs.map((def) => [
            def.displayName,
            { tagSlot: def.tagSlot, fieldType: def.fieldType },
          ])
        )

        for (const filter of validatedData.tagFilters) {
          const current = perKbMap.get(filter.tagName)
          if (!current) {
            if (accessibleKbIds.length > 1) {
              return NextResponse.json(
                {
                  error: `Tag "${filter.tagName}" does not exist in all selected knowledge bases. Search those knowledge bases separately.`,
                },
                { status: 400 }
              )
            }
            continue
          }

          const existing = displayNameToTagDef[filter.tagName]
          if (
            existing &&
            (existing.tagSlot !== current.tagSlot || existing.fieldType !== current.fieldType)
          ) {
            return NextResponse.json(
              {
                error: `Tag "${filter.tagName}" is not mapped consistently across the selected knowledge bases. Search those knowledge bases separately.`,
              },
              { status: 400 }
            )
          }

          displayNameToTagDef[filter.tagName] = current
        }

        logger.debug(`[${requestId}] Loaded tag definitions for KB ${kbId}`, {
          tagCount: tagDefs.length,
        })
      }

      const undefinedTags: string[] = []
      const typeErrors: string[] = []

      for (const filter of validatedData.tagFilters) {
        const tagDef = displayNameToTagDef[filter.tagName]

        if (!tagDef) {
          undefinedTags.push(filter.tagName)
          continue
        }

        const validationError = validateTagValue(
          filter.tagName,
          String(filter.value),
          tagDef.fieldType
        )
        if (validationError) {
          typeErrors.push(validationError)
        }
      }

      if (undefinedTags.length > 0 || typeErrors.length > 0) {
        const errorParts: string[] = []

        if (undefinedTags.length > 0) {
          errorParts.push(buildUndefinedTagsError(undefinedTags))
        }

        if (typeErrors.length > 0) {
          errorParts.push(...typeErrors)
        }

        return NextResponse.json({ error: errorParts.join('\n') }, { status: 400 })
      }

      structuredFilters = validatedData.tagFilters.map((filter) => {
        const tagDef = displayNameToTagDef[filter.tagName]!
        const tagSlot = tagDef.tagSlot
        const fieldType = tagDef.fieldType

        logger.debug(
          `[${requestId}] Structured filter: ${filter.tagName} -> ${tagSlot} (${fieldType}) ${filter.operator}`
        )

        return {
          tagSlot,
          fieldType,
          operator: filter.operator,
          value: filter.value,
          valueTo: filter.valueTo,
        }
      })
    }

    if (accessibleKbIds.length === 0) {
      return NextResponse.json(
        { error: 'Knowledge base not found or access denied' },
        { status: 404 }
      )
    }

    const accessibleKbs = accessChecks
      .filter((ac): ac is KnowledgeBaseAccessResult => Boolean(ac?.hasAccess))
      .map((ac) => ac.knowledgeBase)
    const useReranker = validatedData.rerankerEnabled && Boolean(validatedData.query?.trim())
    const rerankerModel = useReranker ? validatedData.rerankerModel : null

    const hasQuery = validatedData.query && validatedData.query.trim().length > 0
    const workspaceIds = new Set(accessibleKbs.map((kb) => kb.workspaceId ?? null))
    if (hasQuery && workspaceIds.size > 1) {
      return NextResponse.json(
        { error: 'Selected knowledge bases must belong to the same workspace' },
        { status: 400 }
      )
    }
    const workspaceId = accessibleKbs[0]?.workspaceId

    if (workflowId) {
      const authorization = await authorizeWorkflowByWorkspacePermission({
        workflowId: workflowId as string,
        userId,
        action: 'read',
      })
      const workflowWorkspaceId = authorization.workflow?.workspaceId ?? null
      if (
        workflowWorkspaceId &&
        accessChecks.some(
          (accessCheck) =>
            accessCheck?.hasAccess && accessCheck.knowledgeBase?.workspaceId !== workflowWorkspaceId
        )
      ) {
        return NextResponse.json(
          { error: 'Knowledge base does not belong to the workflow workspace' },
          { status: 400 }
        )
      }
    }

    const billingAttribution =
      hasQuery && workspaceId
        ? auth.authType === AuthType.INTERNAL_JWT
          ? requireBillingAttributionHeader(request.headers, {
              actorUserId: userId,
              workspaceId,
            })
          : shouldMeter
            ? await resolveBillingAttribution({
                actorUserId: userId,
                workspaceId,
              })
            : undefined
        : undefined
    const embeddingModels = Array.from(new Set(accessibleKbs.map((kb) => kb.embeddingModel)))
    if (hasQuery && embeddingModels.length > 1) {
      return NextResponse.json(
        {
          error:
            'Selected knowledge bases use different embedding models and cannot be searched together. Search them separately.',
        },
        { status: 400 }
      )
    }
    const queryEmbeddingModel = embeddingModels[0]

    const inaccessibleKbIds = knowledgeBaseIds.filter((id) => !accessibleKbIds.includes(id))

    if (inaccessibleKbIds.length > 0) {
      return NextResponse.json(
        { error: `Knowledge bases not found or access denied: ${inaccessibleKbIds.join(', ')}` },
        { status: 404 }
      )
    }

    /**
     * Gate the workspace payer and actor before hosted embedding cost. Internal
     * workflow tools were gated during preprocessing, and tag-only search is free.
     */
    if (shouldMeter && hasQuery) {
      const usage = billingAttribution
        ? await checkAttributedUsageLimits(billingAttribution)
        : await checkActorUsageLimits(userId)
      if (usage.isExceeded) {
        return NextResponse.json(
          { error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' },
          { status: 402 }
        )
      }
    }

    const modelInputProvenance = await prepareKnowledgeModelInputProvenance({
      headers: request.headers,
      payload: body,
      isInternalRequest: auth.authType === AuthType.INTERNAL_JWT,
      userId,
      workspaceId: workspaceId ?? undefined,
      modelInput: validatedData.query,
    })
    if (!modelInputProvenance.success) {
      return NextResponse.json(
        { error: modelInputProvenance.error },
        { status: modelInputProvenance.status }
      )
    }

    const queryEmbeddingPromise = hasQuery
      ? runWithKnowledgeModelInputProvenance(modelInputProvenance.registry, () =>
          generateSearchEmbedding(validatedData.query!, queryEmbeddingModel, workspaceId)
        )
      : Promise.resolve(null)

    let results: SearchResult[]

    const hasFilters = structuredFilters && structuredFilters.length > 0

    /** Oversample vector results when reranking so the reranker has more to choose from.
     * Cap at 100 to bound Cohere request cost (1 search unit = ≤100 docs). When the caller
     * supplies `rerankerInputCount`, honor it but never let it drop below `topK`
     * (which would defeat the purpose) or exceed 100 (which would split into >1 search units). */
    const rawInputCount = validatedData.rerankerInputCount
    if (useReranker && rawInputCount !== undefined && rawInputCount < validatedData.topK) {
      logger.warn(
        `[${requestId}] rerankerInputCount (${rawInputCount}) is below topK (${validatedData.topK}); raising to topK`
      )
    }
    const candidateTopK = useReranker
      ? rawInputCount !== undefined
        ? Math.min(100, Math.max(validatedData.topK, rawInputCount))
        : Math.min(100, validatedData.topK * 4)
      : validatedData.topK

    if (!hasQuery && hasFilters) {
      results = await executeKnowledgeSearch({
        knowledgeBaseIds: accessibleKbIds,
        topK: validatedData.topK,
        searchMode: validatedData.searchMode,
        structuredFilters,
      })
    } else if (hasQuery) {
      logger.debug(
        `[${requestId}] Executing ${validatedData.searchMode} search`,
        hasFilters ? { filterCount: structuredFilters?.length ?? 0 } : undefined
      )
      const queryVector = JSON.stringify((await queryEmbeddingPromise)?.embedding ?? null)

      results = await executeKnowledgeSearch({
        knowledgeBaseIds: accessibleKbIds,
        topK: candidateTopK,
        searchMode: validatedData.searchMode,
        query: validatedData.query,
        queryVector,
        structuredFilters: hasFilters ? structuredFilters : undefined,
      })
    } else {
      return NextResponse.json(
        {
          error:
            'Please provide either a search query or tag filters to search your knowledge base',
        },
        { status: 400 }
      )
    }

    const resultSecretRegistry =
      modelInputProvenance.registry ??
      new ResolvedSecretTraceRegistry([], {
        userId,
        ...(workspaceId ? { workspaceId } : {}),
      })
    const resultProvenanceSnapshot = await importKnowledgeSearchResultSecretProvenance({
      registry: resultSecretRegistry,
      results,
    })
    if (!resultProvenanceSnapshot.imported) {
      resultSecretRegistry.markIncomplete()
      if (useReranker) {
        return NextResponse.json(
          { error: 'Knowledge result secret provenance is unavailable' },
          { status: 422 }
        )
      }
    }

    /** Optional Cohere rerank pass on top of vector results.
     * `rerankBilled` = Cohere was successfully called (even with 0 results) and we owe the search unit. */
    const rerankedScores = new Map<string, number>()
    let rerankBilled = false
    let rerankIsBYOK = false
    if (useReranker && rerankerModel && results.length > 0) {
      const candidateCount = results.length
      try {
        const { results: ranked, isBYOK } = await runWithKnowledgeModelInputProvenance(
          resultSecretRegistry,
          () =>
            rerank(
              validatedData.query!,
              results.map((r) => ({ id: r.id, text: r.content })),
              {
                model: rerankerModel,
                topN: validatedData.topK,
                workspaceId,
                apiKey: validatedData.rerankerApiKey,
              }
            )
        )
        rerankBilled = true
        rerankIsBYOK = isBYOK
        if (ranked.length === 0) {
          logger.warn(
            `[${requestId}] Reranker returned 0 results; falling back to vector ordering`,
            { model: rerankerModel, candidateCount }
          )
          results = results.slice(0, validatedData.topK)
        } else {
          const idToResult = new Map(results.map((r) => [r.id, r]))
          results = ranked
            .map((r) => idToResult.get(r.item.id))
            .filter((r): r is SearchResult => Boolean(r))
          for (const r of ranked) rerankedScores.set(r.item.id, r.relevanceScore)
          logger.info(`[${requestId}] Reranked ${candidateCount} → ${results.length} results`, {
            model: rerankerModel,
          })
        }
      } catch (error) {
        if (resultSecretRegistry.isPermanentlyIncomplete()) throw error
        logger.warn(`[${requestId}] Reranker failed; falling back to vector ordering`, {
          error: getErrorMessage(error, 'Unknown error'),
          model: rerankerModel,
          candidateCount,
          workspaceId,
        })
        results = results.slice(0, validatedData.topK)
      }
    } else if (useReranker) {
      results = results.slice(0, validatedData.topK)
    }

    let cost = null
    let tokenCount = null
    if (hasQuery) {
      try {
        tokenCount = estimateTokenCount(
          validatedData.query!,
          getEmbeddingModelInfo(queryEmbeddingModel).tokenizerProvider
        )
        // BYOK query embeddings incur no Sim cost, so don't bill (or roll up) them.
        const queryEmbeddingResult = await queryEmbeddingPromise
        if (!queryEmbeddingResult?.isBYOK) {
          cost = calculateCost(queryEmbeddingModel, tokenCount.count, 0, false)
        }
      } catch (error) {
        logger.warn(`[${requestId}] Failed to calculate cost for search query`, {
          error: getErrorMessage(error, 'Unknown error'),
        })
      }
    }

    /** Add Cohere rerank cost (1 search unit per successful call, since we cap candidates ≤100).
     * Bill on every successful API response — Cohere charges even when 0 results are returned. */
    let rerankerCost = 0
    if (rerankBilled && rerankerModel && !rerankIsBYOK) {
      const pricing = getRerankModelPricing(rerankerModel)
      if (pricing) {
        rerankerCost = pricing.perSearchUnit
        if (cost) {
          cost = {
            ...cost,
            input: cost.input + rerankerCost,
            total: cost.total + rerankerCost,
          }
        } else {
          cost = {
            input: rerankerCost,
            output: 0,
            total: rerankerCost,
            pricing: { input: 0, output: 0, updatedAt: pricing.updatedAt },
          }
        }
      } else {
        logger.warn(`[${requestId}] No pricing entry for rerank model ${rerankerModel}`)
      }
    }

    // Record query-embedding + reranker cost for standalone callers (UI, copilot,
    // guardrail RAG). The workflow tool sets skipUsageBilling and rolls the cost
    // up via the executor instead, so this never double-bills; BYOK already
    // resolved to 0 above.
    if (shouldMeter && cost && cost.total > 0) {
      const { recordUsage } = await import('@/lib/billing/core/usage-log')
      try {
        await recordUsage({
          userId,
          workspaceId: workspaceId ?? undefined,
          ...(billingAttribution ? toBillingContext(billingAttribution) : {}),
          entries: [
            {
              category: 'model',
              source: 'knowledge-base',
              description: queryEmbeddingModel,
              cost: cost.total,
              sourceReference: `kb-search:${requestId}`,
            },
          ],
        })
        if (billingAttribution) {
          await checkAndBillPayerOverageThreshold(billingAttribution.billingEntity)
        } else {
          await checkAndBillOverageThreshold(userId)
        }
      } catch (billingError) {
        logger.error(`[${requestId}] Failed to record KB search usage`, { error: billingError })
      }
    }

    const tagDefsResults = await Promise.all(
      accessibleKbIds.map(async (kbId) => {
        try {
          const tagDefs = await getDocumentTagDefinitions(kbId)
          const map: Record<string, string> = {}
          tagDefs.forEach((def) => {
            map[def.tagSlot] = def.displayName
          })
          return { kbId, map }
        } catch (error) {
          logger.warn(`[${requestId}] Failed to fetch tag definitions for display mapping:`, error)
          return { kbId, map: {} as Record<string, string> }
        }
      })
    )
    const tagDefinitionsMap: Record<string, Record<string, string>> = {}
    tagDefsResults.forEach(({ kbId, map }) => {
      tagDefinitionsMap[kbId] = map
    })

    const documentMetadataMap = resultProvenanceSnapshot.documentMetadata

    try {
      PlatformEvents.knowledgeBaseSearched({
        knowledgeBaseId: accessibleKbIds[0],
        resultsCount: results.length,
        workspaceId: workspaceId || undefined,
      })
    } catch {
      // Telemetry should not fail the operation
    }

    const renderedResults = results.map((result) => {
      const kbTagMap = tagDefinitionsMap[result.knowledgeBaseId] || {}
      logger.debug(
        `[${requestId}] Result KB: ${result.knowledgeBaseId}, available mappings:`,
        kbTagMap
      )

      const tags: Record<string, unknown> = {}
      const docMeta = documentMetadataMap[result.documentId]
      ALL_TAG_SLOTS.forEach((slot) => {
        const tagValue = slot.startsWith('tag')
          ? docMeta?.[
              slot as keyof Pick<
                typeof docMeta,
                'tag1' | 'tag2' | 'tag3' | 'tag4' | 'tag5' | 'tag6' | 'tag7'
              >
            ]
          : result[slot]
        if (tagValue !== null && tagValue !== undefined) {
          const displayName = kbTagMap[slot] || slot
          logger.debug(`[${requestId}] Mapping ${slot} -> "${displayName}"`)
          tags[displayName] = tagValue
        }
      })

      const rerankerScore = rerankedScores.get(result.id)
      return {
        documentId: result.documentId,
        documentName: docMeta?.filename || undefined,
        sourceUrl: docMeta?.sourceUrl ?? null,
        content: result.content,
        chunkIndex: result.chunkIndex,
        metadata: tags,
        similarity: hasQuery ? 1 - result.distance : 1,
        ...(rerankerScore !== undefined && { rerankerScore }),
      }
    })

    for (const [documentId, metadata] of Object.entries(documentMetadataMap)) {
      const renderedMetadata = renderedResults
        .filter((result) => result.documentId === documentId)
        .map((result) => ({
          documentName: result.documentName,
          sourceUrl: result.sourceUrl,
          metadata: result.metadata,
        }))
      if (
        renderedMetadata.length > 0 &&
        !(await importDurableSecretProvenance(
          resultSecretRegistry,
          metadata.provenance,
          renderedMetadata
        ))
      ) {
        resultSecretRegistry.markIncomplete()
      }
    }

    const responseBody = {
      success: true,
      data: {
        results: renderedResults,
        query: validatedData.query || '',
        knowledgeBaseIds: accessibleKbIds,
        knowledgeBaseId: accessibleKbIds[0],
        topK: validatedData.topK,
        totalResults: results.length,
        ...(cost
          ? {
              cost: {
                input: cost.input,
                output: cost.output,
                total: cost.total,
                tokens: {
                  prompt: tokenCount?.count ?? 0,
                  completion: 0,
                  total: tokenCount?.count ?? 0,
                },
                model: queryEmbeddingModel,
                pricing: cost.pricing,
                ...(rerankBilled && !rerankIsBYOK
                  ? { rerankerCost, rerankerModel, rerankerSearchUnits: 1 }
                  : {}),
              },
            }
          : {}),
      },
    }
    return createKnowledgeRegistryResponse({
      request,
      authType: auth.authType,
      body: responseBody,
      registry: resultSecretRegistry,
    })
  } catch (error) {
    return NextResponse.json(
      {
        error: 'Failed to perform vector search',
        message: getErrorMessage(error, 'Unknown error'),
      },
      { status: 500 }
    )
  }
})