simstudioai--sim/apps/sim/tools/knowledge/search.ts

import { DEFAULT_RERANKER_MODEL } from '@/lib/knowledge/reranker-models'
import type { KnowledgeSearchResponse } from '@/tools/knowledge/types'
import { enrichKBTagFiltersSchema } from '@/tools/schema-enrichers'
import { parseTagFilters } from '@/tools/shared/tags'
import type { ToolConfig } from '@/tools/types'

export const knowledgeSearchTool: ToolConfig<any, KnowledgeSearchResponse> = {
  id: 'knowledge_search',
  name: 'Knowledge Search',
  description: 'Search for similar content in a knowledge base using vector similarity',
  version: '1.0.0',

  params: {
    knowledgeBaseId: {
      type: 'string',
      required: true,
      visibility: 'user-or-llm',
      description: 'ID of the knowledge base to search in',
    },
    query: {
      type: 'string',
      required: false,
      visibility: 'user-or-llm',
      description: 'Search query text (optional when using tag filters)',
    },
    topK: {
      type: 'number',
      required: false,
      visibility: 'user-or-llm',
      description: 'Number of most similar results to return (1-100)',
    },
    tagFilters: {
      type: 'array',
      required: false,
      visibility: 'user-or-llm',
      description: 'Array of tag filters with tagName and tagValue properties',
      items: {
        type: 'object',
        properties: {
          tagName: { type: 'string' },
          tagValue: { type: 'string' },
        },
      },
    },
    searchMode: {
      type: 'string',
      required: false,
      visibility: 'user-only',
      description:
        "Retrieval mode: 'vector' (default) uses semantic similarity only, 'hybrid' also runs a full-text leg and fuses both",
    },
    rerankerEnabled: {
      type: 'boolean',
      required: false,
      visibility: 'user-only',
      description: 'Whether to apply Cohere reranking to vector search results',
    },
    rerankerModel: {
      type: 'string',
      required: false,
      visibility: 'user-only',
      description:
        'Cohere rerank model to use (one of: rerank-v4.0-pro, rerank-v4.0-fast, rerank-v3.5)',
    },
    rerankerInputCount: {
      type: 'number',
      required: false,
      visibility: 'user-only',
      description:
        'Number of vector results sent to the Cohere reranker (1–100). Defaults to topK × 4 capped at 100.',
    },
    apiKey: {
      type: 'string',
      required: false,
      visibility: 'user-only',
      description: 'Cohere API key for reranker (self-hosted deployments only)',
    },
  },

  schemaEnrichment: {
    tagFilters: {
      dependsOn: 'knowledgeBaseId',
      enrichSchema: enrichKBTagFiltersSchema,
    },
  },

  request: {
    url: () => '/api/knowledge/search',
    method: 'POST',
    modelInput: {
      mode: 'private-provenance',
      inputPaths: () => [['query']],
    },
    secretProvenance: { response: { incomplete: 'reject' } },
    headers: () => ({
      'Content-Type': 'application/json',
    }),
    body: (params) => {
      const workflowId = params._context?.workflowId

      // Use single knowledge base ID
      const knowledgeBaseIds = [params.knowledgeBaseId]

      // Parse tag filters from various formats (array, JSON string)
      const structuredFilters = parseTagFilters(params.tagFilters)

      const rerankerEnabled = params.rerankerEnabled === true || params.rerankerEnabled === 'true'
      const rerankerModel =
        typeof params.rerankerModel === 'string' && params.rerankerModel.length > 0
          ? params.rerankerModel
          : DEFAULT_RERANKER_MODEL
      const rerankerApiKey =
        typeof params.apiKey === 'string' && params.apiKey.length > 0 ? params.apiKey : undefined
      const rawInputCount =
        params.rerankerInputCount !== undefined &&
        params.rerankerInputCount !== null &&
        params.rerankerInputCount !== ''
          ? Number(params.rerankerInputCount)
          : Number.NaN
      const rerankerInputCount = Number.isFinite(rawInputCount)
        ? Math.max(1, Math.min(100, Math.floor(rawInputCount)))
        : undefined

      const requestBody = {
        knowledgeBaseIds,
        query: params.query,
        topK: params.topK ? Math.max(1, Math.min(100, Number(params.topK))) : 10,
        ...(structuredFilters.length > 0 && { tagFilters: structuredFilters }),
        ...(params.searchMode === 'hybrid' && { searchMode: 'hybrid' }),
        ...(rerankerEnabled && {
          rerankerEnabled: true,
          rerankerModel,
          ...(rerankerInputCount !== undefined && { rerankerInputCount }),
          ...(rerankerApiKey && { rerankerApiKey }),
        }),
        ...(workflowId && { workflowId }),
        // The executor rolls this search's cost up at workflow completion, so
        // tell the route not to also meter it (avoids double-billing).
        skipUsageBilling: true,
      }

      return requestBody
    },
  },
  transformResponse: async (response): Promise<KnowledgeSearchResponse> => {
    const result = await response.json()
    const data = result.data || result

    // Restructure cost: extract tokens/model to top level for logging
    let costFields: Record<string, unknown> = {}
    if (data.cost && typeof data.cost === 'object') {
      const {
        tokens,
        model,
        input,
        output: outputCost,
        total,
        rerankerCost,
        rerankerModel,
        rerankerSearchUnits,
      } = data.cost
      costFields = {
        cost: {
          input,
          output: outputCost,
          total,
          ...(typeof rerankerCost === 'number' && { rerankerCost }),
          ...(typeof rerankerModel === 'string' && { rerankerModel }),
          ...(typeof rerankerSearchUnits === 'number' && { rerankerSearchUnits }),
        },
        ...(tokens && { tokens }),
        ...(model && { model }),
      }
    }

    return {
      success: true,
      output: {
        results: data.results || [],
        query: data.query,
        totalResults: data.totalResults || 0,
        ...costFields,
      },
    }
  },

  outputs: {
    results: {
      type: 'array',
      description: 'Array of search results from the knowledge base',
      items: {
        type: 'object',
        properties: {
          documentId: { type: 'string', description: 'Document ID' },
          documentName: { type: 'string', description: 'Document name' },
          sourceUrl: {
            type: 'string',
            nullable: true,
            description:
              'URL to the original source document (e.g., Confluence page, Google Doc, Notion page). Null for documents without an external source.',
          },
          content: { type: 'string', description: 'Content of the result' },
          chunkIndex: { type: 'number', description: 'Index of the chunk within the document' },
          similarity: { type: 'number', description: 'Similarity score of the result' },
          metadata: { type: 'object', description: 'Metadata of the result, including tags' },
        },
      },
    },
    query: {
      type: 'string',
      description: 'The search query that was executed',
    },
    totalResults: {
      type: 'number',
      description: 'Total number of results found',
    },
    cost: {
      type: 'object',
      description: 'Cost information for the search operation',
      optional: true,
    },
  },
}