soulcraft/brainy
1
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions

- Add logging configuration to BrainyData for better control over verbose output.

Browse source 
- Refactor embedding functions to support customizable verbosity through new helper methods (`getDefaultEmbeddingFunction`, `getDefaultBatchEmbeddingFunction`).
- Ensure metadata initialization with default values when null during search results.
This commit is contained in:
David Snelling 2025-07-18 10:40:37 -07:00
parent 51892e46b8
commit d3e6dbd898
2 changed files with 152 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

52
src/brainyData.ts
View file

@ -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<T = any> implements BrainyDataInterface<T> {
@ -142,6 +156,7 @@ export class BrainyData<T = any> implements BrainyDataInterface<T> {
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<T = any> implements BrainyDataInterface<T> {
// 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<T = any> implements BrainyDataInterface<T> {
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<T = any> implements BrainyDataInterface<T> {
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
})
}

135
src/utils/embedding.ts
View file

@ -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<Vector> => {
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<Vector[]> = async (dataArray: string[]): Promise<Vector[]> => {
try {
// Initialize the model if it hasn't been initialized yet
if (!sharedBatchModelInitialized) {
await sharedBatchModel.init()
sharedBatchModelInitialized = true
}
) => Promise<Vector[]> {
// 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<Vector[]> => {
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<Vector[]> {
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