simstudioai--sim/apps/sim/lib/copilot/tools/server/docs/search-documentation.ts

import { db } from '@sim/db'
import { docsEmbeddings } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { sql } from 'drizzle-orm'
import { SearchDocumentation } from '@/lib/copilot/generated/tool-catalog-v1'
import type { BaseServerTool } from '@/lib/copilot/tools/server/base-tool'
import { generateSearchEmbedding } from '@/lib/knowledge/embeddings'

interface DocsSearchParams {
  query: string
  topK?: number
  threshold?: number
}

const DEFAULT_DOCS_SIMILARITY_THRESHOLD = 0.3

export const searchDocumentationServerTool: BaseServerTool<DocsSearchParams, any> = {
  name: SearchDocumentation.id,
  async execute(params: DocsSearchParams): Promise<any> {
    const logger = createLogger('SearchDocumentationServerTool')
    const { query, topK = 10, threshold } = params
    if (!query || typeof query !== 'string') throw new Error('query is required')

    logger.info('Executing docs search', { query, topK })

    const similarityThreshold = threshold ?? DEFAULT_DOCS_SIMILARITY_THRESHOLD

    const { embedding: queryEmbedding } = await generateSearchEmbedding(query)
    if (!queryEmbedding || queryEmbedding.length === 0) {
      return { results: [], query, totalResults: 0 }
    }

    const results = await db
      .select({
        chunkId: docsEmbeddings.chunkId,
        chunkText: docsEmbeddings.chunkText,
        sourceDocument: docsEmbeddings.sourceDocument,
        sourceLink: docsEmbeddings.sourceLink,
        headerText: docsEmbeddings.headerText,
        headerLevel: docsEmbeddings.headerLevel,
        similarity: sql<number>`1 - (${docsEmbeddings.embedding} <=> ${JSON.stringify(queryEmbedding)}::vector)`,
      })
      .from(docsEmbeddings)
      .orderBy(sql`${docsEmbeddings.embedding} <=> ${JSON.stringify(queryEmbedding)}::vector`)
      .limit(topK)

    const filteredResults = results.filter((r) => r.similarity >= similarityThreshold)
    const documentationResults = filteredResults.map((r, idx) => ({
      id: idx + 1,
      title: String(r.headerText || 'Untitled Section'),
      url: String(r.sourceLink || '#'),
      content: String(r.chunkText || ''),
      similarity: r.similarity,
    }))

    logger.info('Docs search complete', { count: documentationResults.length })
    return { results: documentationResults, query, totalResults: documentationResults.length }
  },
}