chore: import upstream snapshot with attribution

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

@@ -0,0 +1,546 @@
+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 }
+    )
+  }
+})

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

@@ -0,0 +1,408 @@
+/**
+ * Tests for knowledge search utility functions
+ * Focuses on testing core functionality with simplified mocking
+ *
+ * @vitest-environment node
+ */
+import { createEnvMock } from '@sim/testing'
+import { mockNextFetchResponse } from '@sim/testing/mocks'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+vi.mock('drizzle-orm')
+vi.mock('@/lib/knowledge/documents/utils', () => ({
+  retryWithExponentialBackoff: (fn: any) => fn(),
+}))
+
+vi.mock('@/lib/core/config/env', () => createEnvMock())
+
+import {
+  generateSearchEmbedding,
+  handleTagAndVectorSearch,
+  handleTagOnlySearch,
+  handleVectorOnlySearch,
+} from '@/app/api/knowledge/search/utils'
+
+describe('Knowledge Search Utils', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  describe('handleTagOnlySearch', () => {
+    it('should throw error when no filters provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [],
+      }
+
+      await expect(handleTagOnlySearch(params)).rejects.toThrow(
+        'Tag filters are required for tag-only search'
+      )
+    })
+
+    it('should accept valid parameters for tag-only search', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'api' }],
+      }
+
+      // This test validates the function accepts the right parameters
+      // The actual database interaction is tested via route tests
+      expect(params.knowledgeBaseIds).toEqual(['kb-123'])
+      expect(params.topK).toBe(10)
+      expect(params.structuredFilters).toHaveLength(1)
+    })
+  })
+
+  describe('handleVectorOnlySearch', () => {
+    it('should throw error when queryVector not provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        distanceThreshold: 0.8,
+      }
+
+      await expect(handleVectorOnlySearch(params)).rejects.toThrow(
+        'Query vector and distance threshold are required for vector-only search'
+      )
+    })
+
+    it('should throw error when distanceThreshold not provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      }
+
+      await expect(handleVectorOnlySearch(params)).rejects.toThrow(
+        'Query vector and distance threshold are required for vector-only search'
+      )
+    })
+
+    it('should accept valid parameters for vector-only search', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        distanceThreshold: 0.8,
+      }
+
+      // This test validates the function accepts the right parameters
+      expect(params.knowledgeBaseIds).toEqual(['kb-123'])
+      expect(params.topK).toBe(10)
+      expect(params.queryVector).toBe(JSON.stringify([0.1, 0.2, 0.3]))
+      expect(params.distanceThreshold).toBe(0.8)
+    })
+  })
+
+  describe('handleTagAndVectorSearch', () => {
+    it('should throw error when no filters provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [],
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        distanceThreshold: 0.8,
+      }
+
+      await expect(handleTagAndVectorSearch(params)).rejects.toThrow(
+        'Tag filters are required for tag and vector search'
+      )
+    })
+
+    it('should throw error when queryVector not provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'api' }],
+        distanceThreshold: 0.8,
+      }
+
+      await expect(handleTagAndVectorSearch(params)).rejects.toThrow(
+        'Query vector and distance threshold are required for tag and vector search'
+      )
+    })
+
+    it('should throw error when distanceThreshold not provided', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'api' }],
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      }
+
+      await expect(handleTagAndVectorSearch(params)).rejects.toThrow(
+        'Query vector and distance threshold are required for tag and vector search'
+      )
+    })
+
+    it('should accept valid parameters for tag and vector search', async () => {
+      const params = {
+        knowledgeBaseIds: ['kb-123'],
+        topK: 10,
+        structuredFilters: [{ tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'api' }],
+        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        distanceThreshold: 0.8,
+      }
+
+      // This test validates the function accepts the right parameters
+      expect(params.knowledgeBaseIds).toEqual(['kb-123'])
+      expect(params.topK).toBe(10)
+      expect(params.structuredFilters).toHaveLength(1)
+      expect(params.queryVector).toBe(JSON.stringify([0.1, 0.2, 0.3]))
+      expect(params.distanceThreshold).toBe(0.8)
+    })
+  })
+
+  describe('generateSearchEmbedding', () => {
+    it('should use Azure OpenAI when KB-specific config is provided', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      const result = await generateSearchEmbedding('test query')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        'https://test.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-12-01-preview',
+        expect.objectContaining({
+          headers: expect.objectContaining({
+            'api-key': 'test-azure-key',
+          }),
+        })
+      )
+      expect(result.embedding).toEqual([0.1, 0.2, 0.3])
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should fallback to OpenAI when no KB Azure config provided', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      const result = await generateSearchEmbedding('test query')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        'https://api.openai.com/v1/embeddings',
+        expect.objectContaining({
+          headers: expect.objectContaining({
+            Authorization: 'Bearer test-openai-key',
+          }),
+        })
+      )
+      expect(result.embedding).toEqual([0.1, 0.2, 0.3])
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('falls back to OpenAI when AZURE_OPENAI_API_VERSION is not set', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        KB_OPENAI_MODEL_NAME: 'custom-embedding-model',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      await generateSearchEmbedding('test query')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        'https://api.openai.com/v1/embeddings',
+        expect.any(Object)
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should use custom model name when provided in Azure config', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'custom-embedding-model',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      await generateSearchEmbedding('test query', 'text-embedding-3-small')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        'https://test.openai.azure.com/openai/deployments/custom-embedding-model/embeddings?api-version=2024-12-01-preview',
+        expect.any(Object)
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should throw error when no API configuration provided', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow(
+        'OPENAI_API_KEY is not configured'
+      )
+    })
+
+    it('should handle Azure OpenAI API errors properly', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+      })
+
+      mockNextFetchResponse({
+        ok: false,
+        status: 404,
+        statusText: 'Not Found',
+        text: 'Deployment not found',
+      })
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should handle OpenAI API errors properly', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        ok: false,
+        status: 429,
+        statusText: 'Too Many Requests',
+        text: 'Rate limit exceeded',
+      })
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should include correct request body for Azure OpenAI', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      await generateSearchEmbedding('test query')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          body: JSON.stringify({
+            input: ['test query'],
+            encoding_format: 'float',
+            dimensions: 1536,
+          }),
+        })
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should include correct request body for OpenAI', async () => {
+      const { env } = await import('@/lib/core/config/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      mockNextFetchResponse({
+        json: {
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+          usage: { prompt_tokens: 1, total_tokens: 1 },
+        },
+      })
+
+      await generateSearchEmbedding('test query', 'text-embedding-3-small')
+
+      expect(vi.mocked(fetch)).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          body: JSON.stringify({
+            input: ['test query'],
+            model: 'text-embedding-3-small',
+            encoding_format: 'float',
+            dimensions: 1536,
+          }),
+        })
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+  })
+
+  describe('getDocumentMetadataByIds', () => {
+    it('should handle empty input gracefully', async () => {
+      const { getDocumentMetadataByIds } = await import('./utils')
+
+      const result = await getDocumentMetadataByIds([])
+
+      expect(result).toEqual({})
+    })
+  })
+})

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

@@ -0,0 +1,539 @@
+import { db } from '@sim/db'
+import { document, embedding } from '@sim/db/schema'
+import { and, eq, inArray, isNull, sql } from 'drizzle-orm'
+import type { StructuredFilter } from '@/lib/knowledge/types'
+
+export interface DocumentMetadata {
+  filename: string
+  sourceUrl: string | null
+}
+
+/**
+ * Batch-fetch display metadata for documents referenced by search results.
+ * Excludes documents that are user-excluded, archived, or soft-deleted —
+ * mirrors the visibility filters applied inside the search SQL itself, so
+ * the lookup will never surface metadata for a row a caller could not have
+ * legitimately matched. Returns a map keyed by document id; missing ids
+ * indicate the document is no longer visible and should be skipped.
+ */
+export async function getDocumentMetadataByIds(
+  documentIds: string[]
+): Promise<Record<string, DocumentMetadata>> {
+  if (documentIds.length === 0) {
+    return {}
+  }
+
+  const uniqueIds = [...new Set(documentIds)]
+  const documents = await db
+    .select({
+      id: document.id,
+      filename: document.filename,
+      sourceUrl: document.sourceUrl,
+    })
+    .from(document)
+    .where(
+      and(
+        inArray(document.id, uniqueIds),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt)
+      )
+    )
+
+  const map: Record<string, DocumentMetadata> = {}
+  documents.forEach((doc) => {
+    map[doc.id] = { filename: doc.filename, sourceUrl: doc.sourceUrl ?? null }
+  })
+
+  return map
+}
+
+export interface SearchResult {
+  id: string
+  content: string
+  documentId: string
+  chunkIndex: number
+  // Text tags
+  tag1: string | null
+  tag2: string | null
+  tag3: string | null
+  tag4: string | null
+  tag5: string | null
+  tag6: string | null
+  tag7: string | null
+  // Number tags (5 slots)
+  number1: number | null
+  number2: number | null
+  number3: number | null
+  number4: number | null
+  number5: number | null
+  // Date tags (2 slots)
+  date1: Date | null
+  date2: Date | null
+  // Boolean tags (3 slots)
+  boolean1: boolean | null
+  boolean2: boolean | null
+  boolean3: boolean | null
+  distance: number
+  knowledgeBaseId: string
+}
+
+export interface SearchParams {
+  knowledgeBaseIds: string[]
+  topK: number
+  structuredFilters?: StructuredFilter[]
+  queryVector?: string
+  distanceThreshold?: number
+}
+
+// Use shared embedding utility
+export { generateSearchEmbedding } from '@/lib/knowledge/embeddings'
+
+/** All valid tag slot keys */
+const TAG_SLOT_KEYS = [
+  // Text tags (7 slots)
+  'tag1',
+  'tag2',
+  'tag3',
+  'tag4',
+  'tag5',
+  'tag6',
+  'tag7',
+  // Number tags (5 slots)
+  'number1',
+  'number2',
+  'number3',
+  'number4',
+  'number5',
+  // Date tags (2 slots)
+  'date1',
+  'date2',
+  // Boolean tags (3 slots)
+  'boolean1',
+  'boolean2',
+  'boolean3',
+] as const
+
+type TagSlotKey = (typeof TAG_SLOT_KEYS)[number]
+
+function isTagSlotKey(key: string): key is TagSlotKey {
+  return TAG_SLOT_KEYS.includes(key as TagSlotKey)
+}
+
+/** Common fields selected for search results */
+const getSearchResultFields = (distanceExpr: any) => ({
+  id: embedding.id,
+  content: embedding.content,
+  documentId: embedding.documentId,
+  chunkIndex: embedding.chunkIndex,
+  // Text tags
+  tag1: embedding.tag1,
+  tag2: embedding.tag2,
+  tag3: embedding.tag3,
+  tag4: embedding.tag4,
+  tag5: embedding.tag5,
+  tag6: embedding.tag6,
+  tag7: embedding.tag7,
+  // Number tags (5 slots)
+  number1: embedding.number1,
+  number2: embedding.number2,
+  number3: embedding.number3,
+  number4: embedding.number4,
+  number5: embedding.number5,
+  // Date tags (2 slots)
+  date1: embedding.date1,
+  date2: embedding.date2,
+  // Boolean tags (3 slots)
+  boolean1: embedding.boolean1,
+  boolean2: embedding.boolean2,
+  boolean3: embedding.boolean3,
+  distance: distanceExpr,
+  knowledgeBaseId: embedding.knowledgeBaseId,
+})
+
+/**
+ * Build a single SQL condition for a filter
+ */
+function buildFilterCondition(filter: StructuredFilter, embeddingTable: any) {
+  const { tagSlot, fieldType, operator, value, valueTo } = filter
+
+  if (!isTagSlotKey(tagSlot)) {
+    return null
+  }
+
+  const column = embeddingTable[tagSlot]
+  if (!column) return null
+
+  // Handle text operators
+  if (fieldType === 'text') {
+    const stringValue = String(value)
+    switch (operator) {
+      case 'eq':
+        return sql`LOWER(${column}) = LOWER(${stringValue})`
+      case 'neq':
+        return sql`LOWER(${column}) != LOWER(${stringValue})`
+      case 'contains':
+        return sql`LOWER(${column}) LIKE LOWER(${`%${stringValue}%`})`
+      case 'not_contains':
+        return sql`LOWER(${column}) NOT LIKE LOWER(${`%${stringValue}%`})`
+      case 'starts_with':
+        return sql`LOWER(${column}) LIKE LOWER(${`${stringValue}%`})`
+      case 'ends_with':
+        return sql`LOWER(${column}) LIKE LOWER(${`%${stringValue}`})`
+      default:
+        return sql`LOWER(${column}) = LOWER(${stringValue})`
+    }
+  }
+
+  // Handle number operators
+  if (fieldType === 'number') {
+    const numValue = typeof value === 'number' ? value : Number.parseFloat(String(value))
+    if (Number.isNaN(numValue)) return null
+
+    switch (operator) {
+      case 'eq':
+        return sql`${column} = ${numValue}`
+      case 'neq':
+        return sql`${column} != ${numValue}`
+      case 'gt':
+        return sql`${column} > ${numValue}`
+      case 'gte':
+        return sql`${column} >= ${numValue}`
+      case 'lt':
+        return sql`${column} < ${numValue}`
+      case 'lte':
+        return sql`${column} <= ${numValue}`
+      case 'between':
+        if (valueTo !== undefined) {
+          const numValueTo =
+            typeof valueTo === 'number' ? valueTo : Number.parseFloat(String(valueTo))
+          if (Number.isNaN(numValueTo)) return sql`${column} = ${numValue}`
+          return sql`${column} >= ${numValue} AND ${column} <= ${numValueTo}`
+        }
+        return sql`${column} = ${numValue}`
+      default:
+        return sql`${column} = ${numValue}`
+    }
+  }
+
+  // Handle date operators - expects YYYY-MM-DD format from frontend
+  if (fieldType === 'date') {
+    const dateStr = String(value)
+    // Validate YYYY-MM-DD format
+    if (!/^\d{4}-\d{2}-\d{2}$/.test(dateStr)) {
+      return null
+    }
+
+    switch (operator) {
+      case 'eq':
+        return sql`${column}::date = ${dateStr}::date`
+      case 'neq':
+        return sql`${column}::date != ${dateStr}::date`
+      case 'gt':
+        return sql`${column}::date > ${dateStr}::date`
+      case 'gte':
+        return sql`${column}::date >= ${dateStr}::date`
+      case 'lt':
+        return sql`${column}::date < ${dateStr}::date`
+      case 'lte':
+        return sql`${column}::date <= ${dateStr}::date`
+      case 'between':
+        if (valueTo !== undefined) {
+          const dateStrTo = String(valueTo)
+          if (!/^\d{4}-\d{2}-\d{2}$/.test(dateStrTo)) {
+            return sql`${column}::date = ${dateStr}::date`
+          }
+          return sql`${column}::date >= ${dateStr}::date AND ${column}::date <= ${dateStrTo}::date`
+        }
+        return sql`${column}::date = ${dateStr}::date`
+      default:
+        return sql`${column}::date = ${dateStr}::date`
+    }
+  }
+
+  // Handle boolean operators
+  if (fieldType === 'boolean') {
+    const boolValue = value === true || value === 'true'
+    switch (operator) {
+      case 'eq':
+        return sql`${column} = ${boolValue}`
+      case 'neq':
+        return sql`${column} != ${boolValue}`
+      default:
+        return sql`${column} = ${boolValue}`
+    }
+  }
+
+  // Fallback to equality
+  return sql`${column} = ${value}`
+}
+
+/**
+ * Build SQL conditions from structured filters with operator support
+ * - Same tag multiple times: OR logic
+ * - Different tags: AND logic
+ */
+function getStructuredTagFilters(filters: StructuredFilter[], embeddingTable: any) {
+  // Group filters by tagSlot
+  const filtersBySlot = new Map<string, StructuredFilter[]>()
+  for (const filter of filters) {
+    const slot = filter.tagSlot
+    if (!filtersBySlot.has(slot)) {
+      filtersBySlot.set(slot, [])
+    }
+    filtersBySlot.get(slot)!.push(filter)
+  }
+
+  // Build conditions: OR within same slot, AND across different slots
+  const conditions: ReturnType<typeof sql>[] = []
+
+  for (const [slot, slotFilters] of filtersBySlot) {
+    const slotConditions = slotFilters
+      .map((f) => buildFilterCondition(f, embeddingTable))
+      .filter((c): c is ReturnType<typeof sql> => c !== null)
+
+    if (slotConditions.length === 0) continue
+
+    if (slotConditions.length === 1) {
+      // Single condition for this slot
+      conditions.push(slotConditions[0])
+    } else {
+      // Multiple conditions for same slot - OR them together
+      conditions.push(sql`(${sql.join(slotConditions, sql` OR `)})`)
+    }
+  }
+
+  return conditions
+}
+
+export function getQueryStrategy(kbCount: number, topK: number) {
+  const useParallel = kbCount > 4 || (kbCount > 2 && topK > 50)
+  const distanceThreshold = kbCount > 3 ? 0.8 : 1.0
+  const parallelLimit = Math.ceil(topK / kbCount) + 5
+
+  return {
+    useParallel,
+    distanceThreshold,
+    parallelLimit,
+    singleQueryOptimized: kbCount <= 2,
+  }
+}
+
+async function executeTagFilterQuery(
+  knowledgeBaseIds: string[],
+  structuredFilters: StructuredFilter[]
+): Promise<{ id: string }[]> {
+  const tagFilterConditions = getStructuredTagFilters(structuredFilters, embedding)
+
+  if (knowledgeBaseIds.length === 1) {
+    return await db
+      .select({ id: embedding.id })
+      .from(embedding)
+      .innerJoin(document, eq(embedding.documentId, document.id))
+      .where(
+        and(
+          eq(embedding.knowledgeBaseId, knowledgeBaseIds[0]),
+          eq(embedding.enabled, true),
+          eq(document.enabled, true),
+          eq(document.processingStatus, 'completed'),
+          eq(document.userExcluded, false),
+          isNull(document.archivedAt),
+          isNull(document.deletedAt),
+          ...tagFilterConditions
+        )
+      )
+  }
+  return await db
+    .select({ id: embedding.id })
+    .from(embedding)
+    .innerJoin(document, eq(embedding.documentId, document.id))
+    .where(
+      and(
+        inArray(embedding.knowledgeBaseId, knowledgeBaseIds),
+        eq(embedding.enabled, true),
+        eq(document.enabled, true),
+        eq(document.processingStatus, 'completed'),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt),
+        ...tagFilterConditions
+      )
+    )
+}
+
+async function executeVectorSearchOnIds(
+  embeddingIds: string[],
+  queryVector: string,
+  topK: number,
+  distanceThreshold: number
+): Promise<SearchResult[]> {
+  if (embeddingIds.length === 0) {
+    return []
+  }
+
+  return await db
+    .select(
+      getSearchResultFields(
+        sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
+      )
+    )
+    .from(embedding)
+    .innerJoin(document, eq(embedding.documentId, document.id))
+    .where(
+      and(
+        inArray(embedding.id, embeddingIds),
+        eq(document.enabled, true),
+        eq(document.processingStatus, 'completed'),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt),
+        sql`${embedding.embedding} <=> ${queryVector}::vector < ${distanceThreshold}`
+      )
+    )
+    .orderBy(sql`${embedding.embedding} <=> ${queryVector}::vector`)
+    .limit(topK)
+}
+
+export async function handleTagOnlySearch(params: SearchParams): Promise<SearchResult[]> {
+  const { knowledgeBaseIds, topK, structuredFilters } = params
+
+  if (!structuredFilters || structuredFilters.length === 0) {
+    throw new Error('Tag filters are required for tag-only search')
+  }
+
+  const strategy = getQueryStrategy(knowledgeBaseIds.length, topK)
+  const tagFilterConditions = getStructuredTagFilters(structuredFilters, embedding)
+
+  if (strategy.useParallel) {
+    // Parallel approach for many KBs
+    const parallelLimit = Math.ceil(topK / knowledgeBaseIds.length) + 5
+
+    const queryPromises = knowledgeBaseIds.map(async (kbId) => {
+      return await db
+        .select(getSearchResultFields(sql<number>`0`.as('distance')))
+        .from(embedding)
+        .innerJoin(document, eq(embedding.documentId, document.id))
+        .where(
+          and(
+            eq(embedding.knowledgeBaseId, kbId),
+            eq(embedding.enabled, true),
+            eq(document.enabled, true),
+            eq(document.processingStatus, 'completed'),
+            eq(document.userExcluded, false),
+            isNull(document.archivedAt),
+            isNull(document.deletedAt),
+            ...tagFilterConditions
+          )
+        )
+        .limit(parallelLimit)
+    })
+
+    const parallelResults = await Promise.all(queryPromises)
+    return parallelResults.flat().slice(0, topK)
+  }
+  // Single query for fewer KBs
+  return await db
+    .select(getSearchResultFields(sql<number>`0`.as('distance')))
+    .from(embedding)
+    .innerJoin(document, eq(embedding.documentId, document.id))
+    .where(
+      and(
+        inArray(embedding.knowledgeBaseId, knowledgeBaseIds),
+        eq(embedding.enabled, true),
+        eq(document.enabled, true),
+        eq(document.processingStatus, 'completed'),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt),
+        ...tagFilterConditions
+      )
+    )
+    .limit(topK)
+}
+
+export async function handleVectorOnlySearch(params: SearchParams): Promise<SearchResult[]> {
+  const { knowledgeBaseIds, topK, queryVector, distanceThreshold } = params
+
+  if (!queryVector || !distanceThreshold) {
+    throw new Error('Query vector and distance threshold are required for vector-only search')
+  }
+
+  const strategy = getQueryStrategy(knowledgeBaseIds.length, topK)
+
+  const distanceExpr = sql<number>`${embedding.embedding} <=> ${queryVector}::vector`.as('distance')
+
+  if (strategy.useParallel) {
+    // Parallel approach for many KBs
+    const parallelLimit = Math.ceil(topK / knowledgeBaseIds.length) + 5
+
+    const queryPromises = knowledgeBaseIds.map(async (kbId) => {
+      return await db
+        .select(getSearchResultFields(distanceExpr))
+        .from(embedding)
+        .innerJoin(document, eq(embedding.documentId, document.id))
+        .where(
+          and(
+            eq(embedding.knowledgeBaseId, kbId),
+            eq(embedding.enabled, true),
+            eq(document.enabled, true),
+            eq(document.processingStatus, 'completed'),
+            eq(document.userExcluded, false),
+            isNull(document.archivedAt),
+            isNull(document.deletedAt),
+            sql`${embedding.embedding} <=> ${queryVector}::vector < ${distanceThreshold}`
+          )
+        )
+        .orderBy(sql`${embedding.embedding} <=> ${queryVector}::vector`)
+        .limit(parallelLimit)
+    })
+
+    const parallelResults = await Promise.all(queryPromises)
+    const allResults = parallelResults.flat()
+    return allResults.sort((a, b) => a.distance - b.distance).slice(0, topK)
+  }
+  // Single query for fewer KBs
+  return await db
+    .select(getSearchResultFields(distanceExpr))
+    .from(embedding)
+    .innerJoin(document, eq(embedding.documentId, document.id))
+    .where(
+      and(
+        inArray(embedding.knowledgeBaseId, knowledgeBaseIds),
+        eq(embedding.enabled, true),
+        eq(document.enabled, true),
+        eq(document.processingStatus, 'completed'),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt),
+        sql`${embedding.embedding} <=> ${queryVector}::vector < ${distanceThreshold}`
+      )
+    )
+    .orderBy(sql`${embedding.embedding} <=> ${queryVector}::vector`)
+    .limit(topK)
+}
+
+export async function handleTagAndVectorSearch(params: SearchParams): Promise<SearchResult[]> {
+  const { knowledgeBaseIds, topK, structuredFilters, queryVector, distanceThreshold } = params
+
+  if (!structuredFilters || structuredFilters.length === 0) {
+    throw new Error('Tag filters are required for tag and vector search')
+  }
+  if (!queryVector || !distanceThreshold) {
+    throw new Error('Query vector and distance threshold are required for tag and vector search')
+  }
+
+  // Step 1: Filter by tags first
+  const tagFilteredIds = await executeTagFilterQuery(knowledgeBaseIds, structuredFilters)
+
+  if (tagFilteredIds.length === 0) {
+    return []
+  }
+
+  // Step 2: Perform vector search only on tag-filtered results
+  return await executeVectorSearchOnIds(
+    tagFilteredIds.map((r) => r.id),
+    queryVector,
+    topK,
+    distanceThreshold
+  )
+}