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 { PlatformEvents } from '@/lib/core/telemetry'
import { generateRequestId } from '@/lib/core/utils/request'
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { ALL_TAG_SLOTS } from '@/lib/knowledge/constants'
import { getEmbeddingModelInfo } from '@/lib/knowledge/embedding-models'
import { rerank } from '@/lib/knowledge/reranker'
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 {
  generateSearchEmbedding,
  getDocumentMetadataByIds,
  getQueryStrategy,
  handleTagAndVectorSearch,
  handleTagOnlySearch,
  handleVectorOnlySearch,
  type SearchResult,
} from '@/app/api/knowledge/search/utils'
import { checkKnowledgeBaseAccess, type KnowledgeBaseAccessResult } from '@/app/api/knowledge/utils'
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} ${filter.value}`
        )

        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 workspaceId = accessibleKbs[0]?.workspaceId

    const useReranker = validatedData.rerankerEnabled && Boolean(validatedData.query?.trim())
    const rerankerModel = useReranker ? validatedData.rerankerModel : null

    const hasQuery = validatedData.query && validatedData.query.trim().length > 0
    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 actor before incurring hosted embedding cost, unless this is the
    // internal workflow tool (already gated at preprocessing, rolls cost up). Tag-only
    // search is free, so only the query path is gated.
    if (shouldMeter && hasQuery) {
      const usage = await checkActorUsageLimits(userId, workspaceId)
      if (usage.isExceeded) {
        return NextResponse.json(
          { error: usage.message || 'Usage limit exceeded. Please upgrade your plan to continue.' },
          { status: 402 }
        )
      }
    }

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

    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 }
        )
      }
    }

    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 handleTagOnlySearch({
        knowledgeBaseIds: accessibleKbIds,
        topK: validatedData.topK,
        structuredFilters,
      })
    } else if (hasQuery && hasFilters) {
      logger.debug(`[${requestId}] Executing tag + vector search with filters:`, structuredFilters)
      const strategy = getQueryStrategy(accessibleKbIds.length, candidateTopK)
      const queryVector = JSON.stringify((await queryEmbeddingPromise)?.embedding ?? null)

      results = await handleTagAndVectorSearch({
        knowledgeBaseIds: accessibleKbIds,
        topK: candidateTopK,
        structuredFilters,
        queryVector,
        distanceThreshold: strategy.distanceThreshold,
      })
    } else if (hasQuery && !hasFilters) {
      const strategy = getQueryStrategy(accessibleKbIds.length, candidateTopK)
      const queryVector = JSON.stringify((await queryEmbeddingPromise)?.embedding ?? null)

      results = await handleVectorOnlySearch({
        knowledgeBaseIds: accessibleKbIds,
        topK: candidateTopK,
        queryVector,
        distanceThreshold: strategy.distanceThreshold,
      })
    } else {
      return NextResponse.json(
        {
          error:
            'Please provide either a search query or tag filters to search your knowledge base',
        },
        { status: 400 }
      )
    }

    /** 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 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) {
        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 && workspaceId && cost && cost.total > 0) {
      const { recordUsage } = await import('@/lib/billing/core/usage-log')
      await recordUsage({
        userId,
        workspaceId,
        entries: [
          {
            category: 'model',
            source: 'knowledge-base',
            description: queryEmbeddingModel,
            cost: cost.total,
            sourceReference: `kb-search:${requestId}`,
          },
        ],
      }).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 documentIds = results.map((result) => result.documentId)
    const documentMetadataMap = await getDocumentMetadataByIds(documentIds)

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

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

          const tags: Record<string, any> = {}

          ALL_TAG_SLOTS.forEach((slot) => {
            const tagValue = (result as any)[slot]
            if (tagValue !== null && tagValue !== undefined) {
              const displayName = kbTagMap[slot] || slot
              logger.debug(
                `[${requestId}] Mapping ${slot}="${tagValue}" -> "${displayName}"="${tagValue}"`
              )
              tags[displayName] = tagValue
            }
          })

          const rerankerScore = rerankedScores.get(result.id)
          const docMeta = documentMetadataMap[result.documentId]
          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 }),
          }
        }),
        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 }
                  : {}),
              },
            }
          : {}),
      },
    })
  } catch (error) {
    return NextResponse.json(
      {
        error: 'Failed to perform vector search',
        message: getErrorMessage(error, 'Unknown error'),
      },
      { status: 500 }
    )
  }
})