From d3e6dbd89879ee37f86073b8ce0d7d4bb2ddadf6 Mon Sep 17 00:00:00 2001 From: dpsifr Date: Fri, 18 Jul 2025 10:40:37 -0700 Subject: [PATCH] - Add logging configuration to `BrainyData` for better control over verbose output. - Refactor embedding functions to support customizable verbosity through new helper methods (`getDefaultEmbeddingFunction`, `getDefaultBatchEmbeddingFunction`). - Ensure metadata initialization with default values when null during search results. --- src/brainyData.ts | 52 +++++++++++++--- src/utils/embedding.ts | 135 ++++++++++++++++++++++++++++++++--------- 2 files changed, 152 insertions(+), 35 deletions(-) diff --git a/src/brainyData.ts b/src/brainyData.ts index 948775ec..9f7a4bf1 100644 --- a/src/brainyData.ts +++ b/src/brainyData.ts @@ -25,6 +25,8 @@ import { cosineDistance, defaultEmbeddingFunction, defaultBatchEmbeddingFunction, + getDefaultEmbeddingFunction, + getDefaultBatchEmbeddingFunction, euclideanDistance, cleanupWorkerPools } from './utils/index.js' @@ -128,6 +130,18 @@ export interface BrainyDataConfig { */ autoConnect?: boolean } + + /** + * Logging configuration + */ + logging?: { + /** + * Whether to enable verbose logging + * When false, suppresses non-essential log messages like model loading progress + * Default: true + */ + verbose?: boolean + } } export class BrainyData implements BrainyDataInterface { @@ -142,6 +156,7 @@ export class BrainyData implements BrainyDataInterface { private storageConfig: BrainyDataConfig['storage'] = {} private useOptimizedIndex: boolean = false private _dimensions: number + private loggingConfig: BrainyDataConfig['logging'] = { verbose: true } // Remote server properties private remoteServerConfig: BrainyDataConfig['remoteServer'] | null = null @@ -203,9 +218,22 @@ export class BrainyData implements BrainyDataInterface { // Set storage if provided, otherwise it will be initialized in init() this.storage = config.storageAdapter || null - // Set embedding function if provided, otherwise use default - this.embeddingFunction = - config.embeddingFunction || defaultEmbeddingFunction + // Store logging configuration + if (config.logging !== undefined) { + this.loggingConfig = { + ...this.loggingConfig, + ...config.logging + } + } + + // Set embedding function if provided, otherwise create one with the appropriate verbose setting + if (config.embeddingFunction) { + this.embeddingFunction = config.embeddingFunction + } else { + this.embeddingFunction = getDefaultEmbeddingFunction({ + verbose: this.loggingConfig.verbose + }) + } // Set persistent storage request flag this.requestPersistentStorage = @@ -806,13 +834,18 @@ export class BrainyData implements BrainyDataInterface { continue } - const metadata = await this.storage!.getMetadata(id) + let metadata = await this.storage!.getMetadata(id) + + // Initialize metadata to an empty object if it's null + if (metadata === null) { + metadata = {} as T + } searchResults.push({ id, score, vector: noun.vector, - metadata: metadata as T | undefined + metadata: metadata as T }) } @@ -855,13 +888,18 @@ export class BrainyData implements BrainyDataInterface { continue } - const metadata = await this.storage!.getMetadata(id) + let metadata = await this.storage!.getMetadata(id) + + // Initialize metadata to an empty object if it's null + if (metadata === null) { + metadata = {} as T + } searchResults.push({ id, score, vector: noun.vector, - metadata: metadata as T | undefined + metadata: metadata as T }) } diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index a8f2a411..877c18ce 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -19,6 +19,15 @@ export class UniversalSentenceEncoder implements EmbeddingModel { private tf: any = null private use: any = null private backend: string = 'cpu' // Default to CPU + private verbose: boolean = true // Whether to log non-essential messages + + /** + * Create a new UniversalSentenceEncoder instance + * @param options Configuration options + */ + constructor(options: { verbose?: boolean } = {}) { + this.verbose = options.verbose !== undefined ? options.verbose : true + } /** * Add polyfills and patches for TensorFlow.js compatibility @@ -93,14 +102,18 @@ export class UniversalSentenceEncoder implements EmbeddingModel { } /** - * Log message only if not in test environment + * Log message only if verbose mode is enabled or if it's an error + * This helps suppress non-essential log messages */ private logger( level: 'log' | 'warn' | 'error', message: string, ...args: any[] ): void { - console[level](message, ...args) + // Always log errors, but only log other messages if verbose mode is enabled + if (level === 'error' || this.verbose) { + console[level](message, ...args) + } } /** @@ -580,14 +593,20 @@ function isTestEnvironment(): boolean { } /** - * Log message only if not in test environment (standalone version) + * Log message only if not in test environment and verbose mode is enabled (standalone version) + * @param level Log level ('log', 'warn', 'error') + * @param message Message to log + * @param args Additional arguments to log + * @param verbose Whether to log non-essential messages (default: true) */ function logIfNotTest( level: 'log' | 'warn' | 'error', message: string, - ...args: any[] + args: any[] = [], + verbose: boolean = true ): void { - if (!isTestEnvironment()) { + // Always log errors, but only log other messages if verbose mode is enabled + if ((level === 'error' || verbose) && !isTestEnvironment()) { console[level](message, ...args) } } @@ -613,18 +632,31 @@ export function createEmbeddingFunction( * Creates a TensorFlow-based Universal Sentence Encoder embedding function * This is the required embedding function for all text embeddings * Uses a shared model instance for better performance across multiple calls + * @param options Configuration options + * @param options.verbose Whether to log non-essential messages (default: true) */ // Create a single shared instance of the model that persists across all embedding calls -const sharedModel = new UniversalSentenceEncoder() +let sharedModel: UniversalSentenceEncoder | null = null let sharedModelInitialized = false +let sharedModelVerbose = true -export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { +export function createTensorFlowEmbeddingFunction(options: { verbose?: boolean } = {}): EmbeddingFunction { + // Update verbose setting if provided + if (options.verbose !== undefined) { + sharedModelVerbose = options.verbose + } + + // Create the shared model if it doesn't exist yet + if (!sharedModel) { + sharedModel = new UniversalSentenceEncoder({ verbose: sharedModelVerbose }) + } + return async (data: any): Promise => { try { // Initialize the model if it hasn't been initialized yet if (!sharedModelInitialized) { try { - await sharedModel.init() + await sharedModel!.init() sharedModelInitialized = true } catch (initError) { // Reset the flag so we can retry initialization on the next call @@ -633,9 +665,9 @@ export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { } } - return await sharedModel.embed(data) + return await sharedModel!.embed(data) } catch (error) { - logIfNotTest('error', 'Failed to use TensorFlow embedding:', error) + logIfNotTest('error', 'Failed to use TensorFlow embedding:', [error], sharedModelVerbose) throw new Error( `Universal Sentence Encoder is required but failed: ${error}` ) @@ -648,40 +680,87 @@ export function createTensorFlowEmbeddingFunction(): EmbeddingFunction { * Uses UniversalSentenceEncoder for all text embeddings * TensorFlow.js is required for this to work * Uses CPU for compatibility + * @param options Configuration options + * @param options.verbose Whether to log non-essential messages (default: true) */ -export const defaultEmbeddingFunction: EmbeddingFunction = - createTensorFlowEmbeddingFunction() +export function getDefaultEmbeddingFunction(options: { verbose?: boolean } = {}): EmbeddingFunction { + return createTensorFlowEmbeddingFunction(options) +} /** - * Default batch embedding function + * Default embedding function with default options * Uses UniversalSentenceEncoder for all text embeddings * TensorFlow.js is required for this to work + * Uses CPU for compatibility + */ +export const defaultEmbeddingFunction: EmbeddingFunction = getDefaultEmbeddingFunction() + +/** + * Creates a batch embedding function that uses UniversalSentenceEncoder + * TensorFlow.js is required for this to work * Processes all items in a single batch operation * Uses a shared model instance for better performance across multiple calls + * @param options Configuration options + * @param options.verbose Whether to log non-essential messages (default: true) */ // Create a single shared instance of the model that persists across function calls -const sharedBatchModel = new UniversalSentenceEncoder() +let sharedBatchModel: UniversalSentenceEncoder | null = null let sharedBatchModelInitialized = false +let sharedBatchModelVerbose = true -export const defaultBatchEmbeddingFunction: ( +export function createBatchEmbeddingFunction(options: { verbose?: boolean } = {}): ( dataArray: string[] -) => Promise = async (dataArray: string[]): Promise => { - try { - // Initialize the model if it hasn't been initialized yet - if (!sharedBatchModelInitialized) { - await sharedBatchModel.init() - sharedBatchModelInitialized = true - } +) => Promise { + // Update verbose setting if provided + if (options.verbose !== undefined) { + sharedBatchModelVerbose = options.verbose + } + + // Create the shared model if it doesn't exist yet + if (!sharedBatchModel) { + sharedBatchModel = new UniversalSentenceEncoder({ verbose: sharedBatchModelVerbose }) + } + + return async (dataArray: string[]): Promise => { + try { + // Initialize the model if it hasn't been initialized yet + if (!sharedBatchModelInitialized) { + await sharedBatchModel!.init() + sharedBatchModelInitialized = true + } - return await sharedBatchModel.embedBatch(dataArray) - } catch (error) { - logIfNotTest('error', 'Failed to use TensorFlow batch embedding:', error) - throw new Error( - `Universal Sentence Encoder batch embedding failed: ${error}` - ) + return await sharedBatchModel!.embedBatch(dataArray) + } catch (error) { + logIfNotTest('error', 'Failed to use TensorFlow batch embedding:', [error], sharedBatchModelVerbose) + throw new Error( + `Universal Sentence Encoder batch embedding failed: ${error}` + ) + } } } +/** + * Get a batch embedding function with custom options + * Uses UniversalSentenceEncoder for all text embeddings + * TensorFlow.js is required for this to work + * Processes all items in a single batch operation + * @param options Configuration options + * @param options.verbose Whether to log non-essential messages (default: true) + */ +export function getDefaultBatchEmbeddingFunction(options: { verbose?: boolean } = {}): ( + dataArray: string[] +) => Promise { + return createBatchEmbeddingFunction(options) +} + +/** + * Default batch embedding function with default options + * Uses UniversalSentenceEncoder for all text embeddings + * TensorFlow.js is required for this to work + * Processes all items in a single batch operation + */ +export const defaultBatchEmbeddingFunction = getDefaultBatchEmbeddingFunction() + /** * Creates an embedding function that runs in a separate thread * This is a wrapper around createEmbeddingFunction that uses executeInThread