brainy/src/utils/embedding.ts

/**
 * Embedding functions for converting data to vectors
 */

import { EmbeddingFunction, EmbeddingModel, Vector } from '../coreTypes.js'
import { executeInThread } from './workerUtils.js'

/**
 * TensorFlow Universal Sentence Encoder embedding model
 * This model provides high-quality text embeddings using TensorFlow.js
 * The required TensorFlow.js dependencies are automatically installed with this package
 *
 * This implementation will use GPU acceleration via WebGL when available,
 * falling back to CPU processing when GPU is not available or fails to initialize.
 */
export class UniversalSentenceEncoder implements EmbeddingModel {
  private model: any = null
  private initialized = false
  private tf: any = null
  private use: any = null

  /**
   * Initialize the embedding model
   */
  public async init(): Promise<void> {
    try {
      // Save original console.warn
      const originalWarn = console.warn

      // Override console.warn to suppress TensorFlow.js Node.js backend message
      console.warn = function (message?: any, ...optionalParams: any[]) {
        if (
          message &&
          typeof message === 'string' &&
          message.includes(
            'Hi, looks like you are running TensorFlow.js in Node.js'
          )
        ) {
          return // Suppress the specific warning
        }
        originalWarn(message, ...optionalParams)
      }

      // TensorFlow.js will use its default EPSILON value

      // Dynamically import TensorFlow.js core module and backends
      // Use type assertions to tell TypeScript these modules exist
      this.tf = await import('@tensorflow/tfjs-core')

      // Import CPU and WebGL backends
      await import('@tensorflow/tfjs-backend-cpu')

      try {
        // Try to import and use WebGL backend first (GPU)
        await import('@tensorflow/tfjs-backend-webgl')

        // Check if WebGL is available and set it as the backend
        if (
          (await this.tf.findBackend('webgl')) ||
          (await this.tf.ready().then(() => this.tf.findBackend('webgl')))
        ) {
          console.log('Using WebGL backend (GPU acceleration)')
          await this.tf.setBackend('webgl')
        } else {
          console.log('WebGL backend not available, falling back to CPU')
          await this.tf.setBackend('cpu')
        }
      } catch (err) {
        console.warn(
          'WebGL backend failed to initialize, using CPU backend:',
          err
        )
        await this.tf.setBackend('cpu')
      }

      this.use = await import('@tensorflow-models/universal-sentence-encoder')

      // Load the model
      this.model = await this.use.load()
      this.initialized = true

      // Restore original console.warn
      console.warn = originalWarn
    } catch (error) {
      console.error('Failed to initialize Universal Sentence Encoder:', error)
      throw new Error(
        `Failed to initialize Universal Sentence Encoder: ${error}`
      )
    }
  }

  /**
   * Embed text into a vector using Universal Sentence Encoder
   * @param data Text to embed
   */
  public async embed(data: string | string[]): Promise<Vector> {
    if (!this.initialized) {
      await this.init()
    }

    try {
      // Handle different input types
      let textToEmbed: string[]
      if (typeof data === 'string') {
        // Handle empty string case
        if (data.trim() === '') {
          // Return a zero vector of appropriate dimension (512 is the default for USE)
          return new Array(512).fill(0)
        }
        textToEmbed = [data]
      } else if (
        Array.isArray(data) &&
        data.every((item) => typeof item === 'string')
      ) {
        // Handle empty array or array with empty strings
        if (data.length === 0 || data.every((item) => item.trim() === '')) {
          return new Array(512).fill(0)
        }
        // Filter out empty strings
        textToEmbed = data.filter((item) => item.trim() !== '')
        if (textToEmbed.length === 0) {
          return new Array(512).fill(0)
        }
      } else {
        throw new Error(
          'UniversalSentenceEncoder only supports string or string[] data'
        )
      }

      // Get embeddings
      const embeddings = await this.model.embed(textToEmbed)

      // Convert to array and return the first embedding
      const embeddingArray = await embeddings.array()
      return embeddingArray[0]
    } catch (error) {
      console.error(
        'Failed to embed text with Universal Sentence Encoder:',
        error
      )
      throw new Error(
        `Failed to embed text with Universal Sentence Encoder: ${error}`
      )
    }
  }

  /**
   * Dispose of the model resources
   */
  public async dispose(): Promise<void> {
    if (this.model && this.tf) {
      try {
        // Dispose of the model and tensors
        this.model.dispose()
        this.tf.disposeVariables()
        this.initialized = false
      } catch (error) {
        console.error('Failed to dispose Universal Sentence Encoder:', error)
      }
    }
    return Promise.resolve()
  }
}

/**
 * Create an embedding function from an embedding model
 * @param model Embedding model to use
 */
export function createEmbeddingFunction(
  model: EmbeddingModel
): EmbeddingFunction {
  return async (data: any): Promise<Vector> => {
    return await model.embed(data)
  }
}

/**
 * Creates a TensorFlow-based Universal Sentence Encoder embedding function
 * This is the required embedding function for all text embeddings
 */
export function createTensorFlowEmbeddingFunction(): EmbeddingFunction {
  // Create a single shared instance of the model
  const model = new UniversalSentenceEncoder()
  let modelInitialized = false

  return async (data: any): Promise<Vector> => {
    try {
      // Initialize the model if it hasn't been initialized yet
      if (!modelInitialized) {
        await model.init()
        modelInitialized = true
      }

      return await model.embed(data)
    } catch (error) {
      console.error('Failed to use TensorFlow embedding:', error)
      throw new Error(
        `Universal Sentence Encoder is required but failed: ${error}`
      )
    }
  }
}

/**
 * Creates a TensorFlow-based Universal Sentence Encoder embedding function that runs in a separate thread
 * This provides better performance for embedding operations by:
 * 1. Using GPU acceleration via WebGL when available
 * 2. Running in a separate thread to avoid blocking the main thread
 * 3. Falling back to CPU processing when GPU is not available
 *
 * @param options Configuration options
 * @returns An embedding function that runs in a separate thread with GPU acceleration when available
 */
export function createThreadedEmbeddingFunction(
  options: { fallbackToMain?: boolean } = {}
): EmbeddingFunction {
  // Create a standard embedding function to use as fallback
  const standardEmbedding = createTensorFlowEmbeddingFunction()

  // Flag to track if we've fallen back to main thread
  let useFallback = false

  return async (data: any): Promise<Vector> => {
    // If we've already determined that threading doesn't work, use the fallback
    if (useFallback) {
      return standardEmbedding(data)
    }

    try {
      // Function to be executed in a worker thread
      // This must be a regular function (not async) to avoid Promise cloning issues
      const embedInWorker = (inputData: any) => {
        // Return a plain object with the input data
        // All async operations will be performed inside the worker
        return { data: inputData }
      }

      // Worker implementation function that will be stringified and run in the worker
      const workerImplementation = async ({ data }: { data: any }) => {
        try {
          // We need to dynamically import TensorFlow.js core module and USE in the worker
          // Use a variable name that won't conflict with any minified variables

          // TensorFlow.js will use its default EPSILON value

          // Import TensorFlow.js modules
          const tf = await import('@tensorflow/tfjs-core')

          // Import CPU backend in the worker
          await import('@tensorflow/tfjs-backend-cpu')

          try {
            // Try to import and use WebGL backend first (GPU)
            await import('@tensorflow/tfjs-backend-webgl')

            // Check if WebGL is available and set it as the backend
            if (
              (tf.findBackend && (await tf.findBackend('webgl'))) ||
              (tf.ready &&
                (await tf
                  .ready()
                  .then(() => tf.findBackend && tf.findBackend('webgl'))))
            ) {
              console.log('Worker: Using WebGL backend (GPU acceleration)')
              if (tf.setBackend) {
                await tf.setBackend('webgl')
              }
            } else {
              console.log(
                'Worker: WebGL backend not available, falling back to CPU'
              )
              if (tf.setBackend) {
                await tf.setBackend('cpu')
              }
            }
          } catch (err) {
            console.warn(
              'Worker: WebGL backend failed to initialize, using CPU backend:',
              err
            )
            if (tf.setBackend) {
              await tf.setBackend('cpu')
            }
          }

          // Import the Universal Sentence Encoder
          const sentenceEncoderModule = await import(
            '@tensorflow-models/universal-sentence-encoder'
          )

          // Load the model directly from the module
          const model = await sentenceEncoderModule.load()

          // Handle different input types
          let textToEmbed: string[]
          if (typeof data === 'string') {
            if (data.trim() === '') {
              return new Array(512).fill(0)
            }
            textToEmbed = [data]
          } else if (
            Array.isArray(data) &&
            data.every((item) => typeof item === 'string')
          ) {
            if (data.length === 0 || data.every((item) => item.trim() === '')) {
              return new Array(512).fill(0)
            }
            textToEmbed = data.filter((item) => item.trim() !== '')
            if (textToEmbed.length === 0) {
              return new Array(512).fill(0)
            }
          } else {
            throw new Error(
              'UniversalSentenceEncoder only supports string or string[] data'
            )
          }

          // Get embeddings
          const embeddings = await model.embed(textToEmbed)

          // Convert to array and return the first embedding
          const embeddingArray = await embeddings.array()

          // Dispose of the tensor to free memory
          embeddings.dispose()

          return embeddingArray[0]
        } catch (error) {
          console.error('Worker error:', error)
          throw error
        }
      }

      // Execute the embedding function in a separate thread
      // Pass the worker implementation as a string to avoid Promise cloning issues
      return await executeInThread<Vector>(
        workerImplementation.toString(),
        embedInWorker(data)
      )
    } catch (error) {
      // If threading fails and fallback is enabled, use the standard embedding function
      if (options.fallbackToMain) {
        console.warn(
          'Threaded embedding failed, falling back to main thread:',
          error
        )
        useFallback = true
        return standardEmbedding(data)
      }

      // Otherwise, propagate the error
      throw new Error(`Threaded embedding failed: ${error}`)
    }
  }
}

/**
 * Default embedding function
 * Uses UniversalSentenceEncoder for all text embeddings
 * TensorFlow.js is required for this to work
 * Uses GPU acceleration via WebGL when available for optimal performance
 * Uses threading when available to avoid blocking the main thread
 * Falls back to CPU processing when GPU is not available
 */
export const defaultEmbeddingFunction: EmbeddingFunction =
  createThreadedEmbeddingFunction({ fallbackToMain: true })