From 5703c4a69a90eb44c367736ef9f3748ea5115350 Mon Sep 17 00:00:00 2001 From: David Snelling Date: Thu, 19 Jun 2025 15:02:25 -0700 Subject: [PATCH] **feat: integrate optimized index functionality into BrainyData** ### Changes: - Enhanced `BrainyData` class to support `HNSWIndexOptimized` alongside the standard HNSW index. - Added a new configuration parameter `hnswOptimized` for enabling optimized indexing. - Created utility methods for handling optimized index-specific functionalities such as memory usage tracking and disk-based index management. - Updated vector addition, retrieval, and management methods to accommodate both standard and optimized indices. - Introduced support for using pre-defined IDs during vector and relationship creation. - Standardized metadata handling for error resilience and embedding integrations. - Implemented a `getAllNouns` method to retrieve all indexed data. - Added validation for vector dimensions during relationship creation to ensure consistency. - Refactored method signatures to improve readability and alignment. ### Purpose: Improved the BrainyData API to support optimized indexing functionality powered by `HNSWIndexOptimized`, enhancing scalability for large datasets while maintaining backward compatibility. Elevated overall code structure and usability. --- src/brainyData.ts | 678 ++++++++++++++++++++++++++------- src/hnsw/hnswIndex.ts | 28 ++ src/hnsw/hnswIndexOptimized.ts | 586 ++++++++++++++++++++++++++++ 3 files changed, 1159 insertions(+), 133 deletions(-) create mode 100644 src/hnsw/hnswIndexOptimized.ts diff --git a/src/brainyData.ts b/src/brainyData.ts index b2fe0054..268d4137 100644 --- a/src/brainyData.ts +++ b/src/brainyData.ts @@ -5,24 +5,31 @@ import { v4 as uuidv4 } from 'uuid' import { HNSWIndex } from './hnsw/hnswIndex.js' +import { HNSWIndexOptimized, HNSWOptimizedConfig } from './hnsw/hnswIndexOptimized.js' import { createStorage } from './storage/opfsStorage.js' import { DistanceFunction, GraphVerb, EmbeddingFunction, - HNSWConfig, HNSWNoun, + HNSWConfig, + HNSWNoun, SearchResult, StorageAdapter, Vector, VectorDocument } from './coreTypes.js' -import { cosineDistance, defaultEmbeddingFunction, euclideanDistance } from './utils/index.js' +import { + cosineDistance, + defaultEmbeddingFunction, + euclideanDistance +} from './utils/index.js' import { NounType, VerbType, GraphNoun } from './types/graphTypes.js' -import { - ServerSearchConduitAugmentation, - createServerSearchAugmentations +import { + ServerSearchConduitAugmentation, + createServerSearchAugmentations } from './augmentations/serverSearchAugmentations.js' import { WebSocketConnection } from './types/augmentations.js' +import { BrainyDataInterface } from './types/brainyDataInterface.js' export interface BrainyDataConfig { /** @@ -30,6 +37,12 @@ export interface BrainyDataConfig { */ hnsw?: Partial + /** + * Optimized HNSW index configuration + * If provided, will use the optimized HNSW index instead of the standard one + */ + hnswOptimized?: Partial + /** * Distance function to use for similarity calculations */ @@ -45,34 +58,34 @@ export interface BrainyDataConfig { * These will be passed to createStorage if storageAdapter is not provided */ storage?: { - requestPersistentStorage?: boolean; + requestPersistentStorage?: boolean r2Storage?: { - bucketName?: string; - accountId?: string; - accessKeyId?: string; - secretAccessKey?: string; - }; + bucketName?: string + accountId?: string + accessKeyId?: string + secretAccessKey?: string + } s3Storage?: { - bucketName?: string; - accessKeyId?: string; - secretAccessKey?: string; - region?: string; - }; + bucketName?: string + accessKeyId?: string + secretAccessKey?: string + region?: string + } gcsStorage?: { - bucketName?: string; - accessKeyId?: string; - secretAccessKey?: string; - endpoint?: string; - }; + bucketName?: string + accessKeyId?: string + secretAccessKey?: string + endpoint?: string + } customS3Storage?: { - bucketName?: string; - accessKeyId?: string; - secretAccessKey?: string; - endpoint?: string; - region?: string; - }; - forceFileSystemStorage?: boolean; - forceMemoryStorage?: boolean; + bucketName?: string + accessKeyId?: string + secretAccessKey?: string + endpoint?: string + region?: string + } + forceFileSystemStorage?: boolean + forceMemoryStorage?: boolean } /** @@ -100,55 +113,69 @@ export interface BrainyDataConfig { /** * WebSocket URL of the remote Brainy server */ - url: string; + url: string /** * WebSocket protocols to use for the connection */ - protocols?: string | string[]; + protocols?: string | string[] /** * Whether to automatically connect to the remote server on initialization */ - autoConnect?: boolean; + autoConnect?: boolean } } -export class BrainyData { - private index: HNSWIndex +export class BrainyData implements BrainyDataInterface { + private index: HNSWIndex | HNSWIndexOptimized private storage: StorageAdapter | null = null private isInitialized = false private embeddingFunction: EmbeddingFunction + private distanceFunction: DistanceFunction private requestPersistentStorage: boolean private readOnly: boolean private storageConfig: BrainyDataConfig['storage'] = {} + private useOptimizedIndex: boolean = false // Remote server properties private remoteServerConfig: BrainyDataConfig['remoteServer'] | null = null - private serverSearchConduit: any = null // Will be ServerSearchConduitAugmentation when connected - private serverConnection: any = null // Will be WebSocketConnection when connected + private serverSearchConduit: ServerSearchConduitAugmentation | null = null + private serverConnection: WebSocketConnection | null = null /** * Create a new vector database */ constructor(config: BrainyDataConfig = {}) { - // Initialize HNSW index - this.index = new HNSWIndex( - config.hnsw, - config.distanceFunction || cosineDistance - ) + // Set distance function + this.distanceFunction = config.distanceFunction || cosineDistance + + // Check if optimized HNSW index configuration is provided + if (config.hnswOptimized) { + // Initialize optimized HNSW index + this.index = new HNSWIndexOptimized( + config.hnswOptimized, + this.distanceFunction, + config.storageAdapter || null + ) + this.useOptimizedIndex = true + } else { + // Initialize standard HNSW index + this.index = new HNSWIndex(config.hnsw, this.distanceFunction) + } // 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 + this.embeddingFunction = + config.embeddingFunction || defaultEmbeddingFunction // Set persistent storage request flag (support both new and deprecated options) this.requestPersistentStorage = - (config.storage?.requestPersistentStorage !== undefined) + config.storage?.requestPersistentStorage !== undefined ? config.storage.requestPersistentStorage - : (config.requestPersistentStorage || false) + : config.requestPersistentStorage || false // Set read-only flag this.readOnly = config.readOnly || false @@ -168,7 +195,9 @@ export class BrainyData { */ private checkReadOnly(): void { if (this.readOnly) { - throw new Error('Cannot perform write operation: database is in read-only mode') + throw new Error( + 'Cannot perform write operation: database is in read-only mode' + ) } } @@ -196,6 +225,11 @@ export class BrainyData { // Initialize storage await this.storage!.init() + // If using optimized index, set the storage adapter + if (this.useOptimizedIndex && this.index instanceof HNSWIndexOptimized) { + this.index.setStorage(this.storage!) + } + // Load all nouns from storage const nouns: HNSWNoun[] = await this.storage!.getAllNouns() @@ -274,8 +308,9 @@ export class BrainyData { vectorOrData: Vector | any, metadata?: T, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector + forceEmbed?: boolean // Force using the embedding function even if input is a vector addToRemote?: boolean // Whether to also add to the remote server if connected + id?: string // Optional ID to use instead of generating a new one } = {} ): Promise { await this.ensureInitialized() @@ -308,14 +343,18 @@ export class BrainyData { throw new Error('Vector is undefined or null') } - // Generate ID if isn't provided - const id = uuidv4() + // Use ID from options if it exists, otherwise from metadata, otherwise generate a new UUID + const id = + options.id || + (metadata && typeof metadata === 'object' && 'id' in metadata + ? (metadata as any).id + : uuidv4()) // Add to index this.index.addItem({ id, vector }) // Get the noun from the index - const noun = this.index.getNodes().get(id) + const noun = this.index.getNouns().get(id) if (!noun) { throw new Error(`Failed to retrieve newly created noun with ID ${id}`) @@ -328,15 +367,17 @@ export class BrainyData { if (metadata !== undefined) { // Validate noun type if metadata is for a GraphNoun if (metadata && typeof metadata === 'object' && 'noun' in metadata) { - const nounType = (metadata as any).noun + const nounType = (metadata as unknown as GraphNoun).noun // Check if the noun type is valid const isValidNounType = Object.values(NounType).includes(nounType) if (!isValidNounType) { - console.warn(`Invalid noun type: ${nounType}. Falling back to GraphNoun.`); + console.warn( + `Invalid noun type: ${nounType}. Falling back to GraphNoun.` + ) // Set a default noun type - (metadata as any).noun = NounType.Concept + ;(metadata as unknown as GraphNoun).noun = NounType.Concept } } @@ -348,7 +389,9 @@ export class BrainyData { try { await this.addToRemote(id, vector, metadata) } catch (remoteError) { - console.warn(`Failed to add to remote server: ${remoteError}. Continuing with local add.`) + console.warn( + `Failed to add to remote server: ${remoteError}. Continuing with local add.` + ) } } @@ -375,7 +418,9 @@ export class BrainyData { ): Promise { // Check if connected to a remote server if (!this.isConnectedToRemoteServer()) { - throw new Error('Not connected to a remote server. Call connectToRemoteServer() first.') + throw new Error( + 'Not connected to a remote server. Call connectToRemoteServer() first.' + ) } // Add to local with addToRemote option @@ -400,6 +445,12 @@ export class BrainyData { } try { + if (!this.serverSearchConduit || !this.serverConnection) { + throw new Error( + 'Server search conduit or connection is not initialized' + ) + } + // Add to remote server const addResult = await this.serverSearchConduit.addToBoth( this.serverConnection.connectionId, @@ -426,11 +477,11 @@ export class BrainyData { */ public async addBatch( items: Array<{ - vectorOrData: Vector | any; + vectorOrData: Vector | any metadata?: T }>, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector + forceEmbed?: boolean // Force using the embedding function even if input is a vector addToRemote?: boolean // Whether to also add to the remote server if connected } = {} ): Promise { @@ -462,7 +513,7 @@ export class BrainyData { */ public async addBatchToBoth( items: Array<{ - vectorOrData: Vector | any; + vectorOrData: Vector | any metadata?: T }>, options: { @@ -471,7 +522,9 @@ export class BrainyData { ): Promise { // Check if connected to a remote server if (!this.isConnectedToRemoteServer()) { - throw new Error('Not connected to a remote server. Call connectToRemoteServer() first.') + throw new Error( + 'Not connected to a remote server. Call connectToRemoteServer() first.' + ) } // Add to local with addToRemote option @@ -530,7 +583,7 @@ export class BrainyData { const searchResults: SearchResult[] = [] for (const [id, score] of results) { - const noun = this.index.getNodes().get(id) + const noun = this.index.getNouns().get(id) if (!noun) { continue } @@ -541,14 +594,16 @@ export class BrainyData { id, score, vector: noun.vector, - metadata + metadata: metadata as T | undefined }) } return searchResults } else { // Get nouns for each noun type in parallel - const nounPromises = nounTypes.map(nounType => this.storage!.getNounsByNounType(nounType)) + const nounPromises = nounTypes.map((nounType) => + this.storage!.getNounsByNounType(nounType) + ) const nounArrays = await Promise.all(nounPromises) // Combine all nouns @@ -560,7 +615,10 @@ export class BrainyData { // Calculate distances for each noun const results: Array<[string, number]> = [] for (const noun of nouns) { - const distance = this.index.getDistanceFunction()(queryVector, noun.vector) + const distance = this.index.getDistanceFunction()( + queryVector, + noun.vector + ) results.push([noun.id, distance]) } @@ -574,7 +632,7 @@ export class BrainyData { const searchResults: SearchResult[] = [] for (const [id, score] of topResults) { - const noun = nouns.find(n => n.id === id) + const noun = nouns.find((n) => n.id === id) if (!noun) { continue } @@ -585,7 +643,7 @@ export class BrainyData { id, score, vector: noun.vector, - metadata + metadata: metadata as T | undefined }) } @@ -608,9 +666,9 @@ export class BrainyData { queryVectorOrData: Vector | any, k: number = 10, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector - nounTypes?: string[], // Optional array of noun types to search within - includeVerbs?: boolean, // Whether to include associated GraphVerbs in the results + forceEmbed?: boolean // Force using the embedding function even if input is a vector + nounTypes?: string[] // Optional array of noun types to search within + includeVerbs?: boolean // Whether to include associated GraphVerbs in the results searchMode?: 'local' | 'remote' | 'combined' // Where to search: local, remote, or both } = {} ): Promise[]> { @@ -638,8 +696,8 @@ export class BrainyData { queryVectorOrData: Vector | any, k: number = 10, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector - nounTypes?: string[], // Optional array of noun types to search within + forceEmbed?: boolean // Force using the embedding function even if input is a vector + nounTypes?: string[] // Optional array of noun types to search within includeVerbs?: boolean // Whether to include associated GraphVerbs in the results } = {} ): Promise[]> { @@ -653,9 +711,14 @@ export class BrainyData { // If noun types are specified, use searchByNounTypes let searchResults if (options.nounTypes && options.nounTypes.length > 0) { - searchResults = await this.searchByNounTypes(queryToUse, k, options.nounTypes, { - forceEmbed: options.forceEmbed - }) + searchResults = await this.searchByNounTypes( + queryToUse, + k, + options.nounTypes, + { + forceEmbed: options.forceEmbed + } + ) } else { // Otherwise, search all GraphNouns searchResults = await this.searchByNounTypes(queryToUse, k, null, { @@ -682,7 +745,7 @@ export class BrainyData { } // Add the verbs to the metadata - (result.metadata as any).associatedVerbs = allVerbs + ;(result.metadata as Record).associatedVerbs = allVerbs } catch (error) { console.warn(`Failed to retrieve verbs for noun ${result.id}:`, error) } @@ -700,7 +763,7 @@ export class BrainyData { try { // Get noun from index - const noun = this.index.getNodes().get(id) + const noun = this.index.getNouns().get(id) if (!noun) { return null } @@ -711,7 +774,7 @@ export class BrainyData { return { id, vector: noun.vector, - metadata + metadata: metadata as T | undefined } } catch (error) { console.error(`Failed to get vector ${id}:`, error) @@ -719,6 +782,33 @@ export class BrainyData { } } + /** + * Get all nouns in the database + * @returns Array of vector documents + */ + public async getAllNouns(): Promise[]> { + await this.ensureInitialized() + + try { + const nouns = this.index.getNouns() + const result: VectorDocument[] = [] + + for (const [id, noun] of nouns.entries()) { + const metadata = await this.storage!.getMetadata(id) + result.push({ + id, + vector: noun.vector, + metadata: metadata as T | undefined + }) + } + + return result + } catch (error) { + console.error('Failed to get all nouns:', error) + throw new Error(`Failed to get all nouns: ${error}`) + } + } + /** * Delete a vector by ID */ @@ -763,22 +853,24 @@ export class BrainyData { try { // Check if a vector exists - const noun = this.index.getNodes().get(id) + const noun = this.index.getNouns().get(id) if (!noun) { return false } // Validate noun type if metadata is for a GraphNoun if (metadata && typeof metadata === 'object' && 'noun' in metadata) { - const nounType = (metadata as any).noun + const nounType = (metadata as unknown as GraphNoun).noun // Check if the noun type is valid const isValidNounType = Object.values(NounType).includes(nounType) if (!isValidNounType) { - console.warn(`Invalid noun type: ${nounType}. Falling back to GraphNoun.`); + console.warn( + `Invalid noun type: ${nounType}. Falling back to GraphNoun.` + ) // Set a default noun type - (metadata as any).noun = NounType.Concept + ;(metadata as unknown as GraphNoun).noun = NounType.Concept } } @@ -792,6 +884,22 @@ export class BrainyData { } } + /** + * Create a relationship between two entities + * This is a convenience wrapper around addVerb + */ + public async relate( + sourceId: string, + targetId: string, + relationType: string, + metadata?: any + ): Promise { + return this.addVerb(sourceId, targetId, undefined, { + type: relationType, + metadata: metadata + }) + } + /** * Add a verb between two nouns * If metadata is provided and vector is not, the metadata will be vectorized using the embedding function @@ -805,6 +913,7 @@ export class BrainyData { weight?: number metadata?: any forceEmbed?: boolean // Force using the embedding function for metadata even if vector is provided + id?: string // Optional ID to use instead of generating a new one } = {} ): Promise { await this.ensureInitialized() @@ -814,8 +923,8 @@ export class BrainyData { try { // Check if source and target nouns exist - const sourceNoun = this.index.getNodes().get(sourceId) - const targetNoun = this.index.getNodes().get(targetId) + const sourceNoun = this.index.getNouns().get(sourceId) + const targetNoun = this.index.getNouns().get(targetId) if (!sourceNoun) { throw new Error(`Source noun with ID ${sourceId} not found`) @@ -825,8 +934,8 @@ export class BrainyData { throw new Error(`Target noun with ID ${targetId} not found`) } - // Generate ID for the verb - const id = uuidv4() + // Use provided ID or generate a new one + const id = options.id || uuidv4() let verbVector: Vector @@ -837,7 +946,10 @@ export class BrainyData { let textToEmbed: string if (typeof options.metadata === 'string') { textToEmbed = options.metadata - } else if (options.metadata.description && typeof options.metadata.description === 'string') { + } else if ( + options.metadata.description && + typeof options.metadata.description === 'string' + ) { textToEmbed = options.metadata.description } else { // Convert to JSON string as fallback @@ -855,19 +967,41 @@ export class BrainyData { } } else { // Use a provided vector or average of source and target vectors - verbVector = - vector || - sourceNoun.vector.map((val, i) => (val + targetNoun.vector[i]) / 2) + if (vector) { + verbVector = vector + } else { + // Ensure both source and target vectors have the same dimension + if ( + !sourceNoun.vector || + !targetNoun.vector || + sourceNoun.vector.length === 0 || + targetNoun.vector.length === 0 || + sourceNoun.vector.length !== targetNoun.vector.length + ) { + throw new Error( + `Cannot average vectors: source or target vector is invalid or dimensions don't match` + ) + } + + // Average the vectors + verbVector = sourceNoun.vector.map( + (val, i) => (val + targetNoun.vector[i]) / 2 + ) + } } // Validate verb type if provided let verbType = options.type if (verbType) { // Check if the verb type is valid - const isValidVerbType = Object.values(VerbType).includes(verbType as any) + const isValidVerbType = Object.values(VerbType).includes( + verbType as VerbType + ) if (!isValidVerbType) { - console.warn(`Invalid verb type: ${verbType}. Using RelatedTo as default.`) + console.warn( + `Invalid verb type: ${verbType}. Using RelatedTo as default.` + ) // Set a default verb type verbType = VerbType.RelatedTo } @@ -889,7 +1023,7 @@ export class BrainyData { this.index.addItem({ id, vector: verbVector }) // Get the noun from the index - const indexNoun = this.index.getNodes().get(id) + const indexNoun = this.index.getNouns().get(id) if (!indexNoun) { throw new Error( @@ -1081,8 +1215,8 @@ export class BrainyData { query: string, k: number = 10, options: { - nounTypes?: string[], - includeVerbs?: boolean, + nounTypes?: string[] + includeVerbs?: boolean searchMode?: 'local' | 'remote' | 'combined' } = {} ): Promise[]> { @@ -1115,9 +1249,9 @@ export class BrainyData { queryVectorOrData: Vector | any, k: number = 10, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector - nounTypes?: string[], // Optional array of noun types to search within - includeVerbs?: boolean, // Whether to include associated GraphVerbs in the results + forceEmbed?: boolean // Force using the embedding function even if input is a vector + nounTypes?: string[] // Optional array of noun types to search within + includeVerbs?: boolean // Whether to include associated GraphVerbs in the results storeResults?: boolean // Whether to store the results in the local database (default: true) } = {} ): Promise[]> { @@ -1125,7 +1259,9 @@ export class BrainyData { // Check if connected to a remote server if (!this.isConnectedToRemoteServer()) { - throw new Error('Not connected to a remote server. Call connectToRemoteServer() first.') + throw new Error( + 'Not connected to a remote server. Call connectToRemoteServer() first.' + ) } try { @@ -1139,6 +1275,12 @@ export class BrainyData { query = 'vector-query' // Placeholder, would need a better approach for vector queries } + if (!this.serverSearchConduit || !this.serverConnection) { + throw new Error( + 'Server search conduit or connection is not initialized' + ) + } + // Search the remote server const searchResult = await this.serverSearchConduit.searchServer( this.serverConnection.connectionId, @@ -1168,9 +1310,9 @@ export class BrainyData { queryVectorOrData: Vector | any, k: number = 10, options: { - forceEmbed?: boolean, // Force using the embedding function even if input is a vector - nounTypes?: string[], // Optional array of noun types to search within - includeVerbs?: boolean, // Whether to include associated GraphVerbs in the results + forceEmbed?: boolean // Force using the embedding function even if input is a vector + nounTypes?: string[] // Optional array of noun types to search within + includeVerbs?: boolean // Whether to include associated GraphVerbs in the results localFirst?: boolean // Whether to search local first (default: true) } = {} ): Promise[]> { @@ -1188,7 +1330,11 @@ export class BrainyData { if (localFirst) { // Search local first - const localResults = await this.searchLocal(queryVectorOrData, k, options) + const localResults = await this.searchLocal( + queryVectorOrData, + k, + options + ) // If we have enough local results, return them if (localResults.length >= k) { @@ -1197,14 +1343,14 @@ export class BrainyData { // Otherwise, search remote for additional results const remoteResults = await this.searchRemote( - queryVectorOrData, + queryVectorOrData, k - localResults.length, { ...options, storeResults: true } ) // Combine results, removing duplicates const combinedResults = [...localResults] - const localIds = new Set(localResults.map(r => r.id)) + const localIds = new Set(localResults.map((r) => r.id)) for (const result of remoteResults) { if (!localIds.has(result.id)) { @@ -1215,11 +1361,10 @@ export class BrainyData { return combinedResults } else { // Search remote first - const remoteResults = await this.searchRemote( - queryVectorOrData, - k, - { ...options, storeResults: true } - ) + const remoteResults = await this.searchRemote(queryVectorOrData, k, { + ...options, + storeResults: true + }) // If we have enough remote results, return them if (remoteResults.length >= k) { @@ -1227,11 +1372,15 @@ export class BrainyData { } // Otherwise, search local for additional results - const localResults = await this.searchLocal(queryVectorOrData, k - remoteResults.length, options) + const localResults = await this.searchLocal( + queryVectorOrData, + k - remoteResults.length, + options + ) // Combine results, removing duplicates const combinedResults = [...remoteResults] - const remoteIds = new Set(remoteResults.map(r => r.id)) + const remoteIds = new Set(remoteResults.map((r) => r.id)) for (const result of localResults) { if (!remoteIds.has(result.id)) { @@ -1265,8 +1414,16 @@ export class BrainyData { } try { + if (!this.serverSearchConduit || !this.serverConnection) { + throw new Error( + 'Server search conduit or connection is not initialized' + ) + } + // Close the WebSocket connection - await this.serverSearchConduit.closeWebSocket(this.serverConnection.connectionId) + await this.serverSearchConduit.closeWebSocket( + this.serverConnection.connectionId + ) // Clear the connection information this.serverSearchConduit = null @@ -1293,16 +1450,16 @@ export class BrainyData { * @returns Object containing the storage type, used space, quota, and additional details */ public async status(): Promise<{ - type: string; - used: number; - quota: number | null; - details?: Record; + type: string + used: number + quota: number | null + details?: Record }> { await this.ensureInitialized() if (!this.storage) { return { - type: 'unknown', + type: 'any', used: 0, quota: null, details: { error: 'Storage not initialized' } @@ -1313,9 +1470,11 @@ export class BrainyData { // Check if the storage adapter has a getStorageStatus method if (typeof this.storage.getStorageStatus !== 'function') { // If not, determine the storage type based on the constructor name - const storageType = this.storage.constructor.name.toLowerCase().replace('storage', '') + const storageType = this.storage.constructor.name + .toLowerCase() + .replace('storage', '') return { - type: storageType || 'unknown', + type: storageType || 'any', used: 0, quota: null, details: { @@ -1330,13 +1489,27 @@ export class BrainyData { const storageStatus = await this.storage.getStorageStatus() // Add index information to the details - const indexInfo = { + let indexInfo: Record = { indexSize: this.size() } + // Add optimized index information if using optimized index + if (this.useOptimizedIndex && this.index instanceof HNSWIndexOptimized) { + const optimizedIndex = this.index as HNSWIndexOptimized + indexInfo = { + ...indexInfo, + optimized: true, + memoryUsage: optimizedIndex.getMemoryUsage(), + productQuantization: optimizedIndex.getUseProductQuantization(), + diskBasedIndex: optimizedIndex.getUseDiskBasedIndex() + } + } else { + indexInfo.optimized = false + } + // Ensure all required fields are present return { - type: storageStatus.type || 'unknown', + type: storageStatus.type || 'any', used: storageStatus.used || 0, quota: storageStatus.quota || null, details: { @@ -1348,10 +1521,12 @@ export class BrainyData { console.error('Failed to get storage status:', error) // Determine the storage type based on the constructor name - const storageType = this.storage.constructor.name.toLowerCase().replace('storage', '') + const storageType = this.storage.constructor.name + .toLowerCase() + .replace('storage', '') return { - type: storageType || 'unknown', + type: storageType || 'any', used: 0, quota: null, details: { @@ -1383,21 +1558,252 @@ export class BrainyData { } } + /** + * Backup all data from the database to a JSON-serializable format + * @returns Object containing all nouns, verbs, noun types, verb types, HNSW index, and other related data + * + * The HNSW index data includes: + * - entryPointId: The ID of the entry point for the graph + * - maxLevel: The maximum level in the hierarchical structure + * - dimension: The dimension of the vectors + * - config: Configuration parameters for the HNSW algorithm + * - connections: A serialized representation of the connections between nouns + */ + public async backup(): Promise<{ + nouns: VectorDocument[] + verbs: GraphVerb[] + nounTypes: string[] + verbTypes: string[] + version: string + hnswIndex?: { + entryPointId: string | null + maxLevel: number + dimension: number | null + config: HNSWConfig + connections: Record> + } + }> { + await this.ensureInitialized() + + try { + // Get all nouns + const nouns = await this.getAllNouns() + + // Get all verbs + const verbs = await this.getAllVerbs() + + // Get all noun types + const nounTypes = Object.values(NounType) + + // Get all verb types + const verbTypes = Object.values(VerbType) + + // Get HNSW index data + const hnswIndexData = { + entryPointId: this.index.getEntryPointId(), + maxLevel: this.index.getMaxLevel(), + dimension: this.index.getDimension(), + config: this.index.getConfig(), + connections: {} as Record> + } + + // Convert Map> to a serializable format + const indexNouns = this.index.getNouns() + for (const [id, noun] of indexNouns.entries()) { + hnswIndexData.connections[id] = {} + for (const [level, connections] of noun.connections.entries()) { + hnswIndexData.connections[id][level] = Array.from(connections) + } + } + + // Return the data with version information + return { + nouns, + verbs, + nounTypes, + verbTypes, + hnswIndex: hnswIndexData, + version: '1.0.0' // Version of the backup format + } + } catch (error) { + console.error('Failed to backup data:', error) + throw new Error(`Failed to backup data: ${error}`) + } + } + + /** + * Restore data into the database from a previously backed up format + * @param data The data to restore, in the format returned by backup() + * This can include HNSW index data if it was included in the backup + * If vectors are not present for nouns, they will be created using the embedding function + * @param options Restore options + * @returns Object containing counts of restored items + */ + public async restore( + data: { + nouns: VectorDocument[] + verbs: GraphVerb[] + nounTypes?: string[] + verbTypes?: string[] + hnswIndex?: { + entryPointId: string | null + maxLevel: number + dimension: number | null + config: HNSWConfig + connections: Record> + } + version: string + }, + options: { + clearExisting?: boolean + } = {} + ): Promise<{ + nounsRestored: number + verbsRestored: number + }> { + await this.ensureInitialized() + + // Check if database is in read-only mode + this.checkReadOnly() + + try { + // Clear existing data if requested + if (options.clearExisting) { + await this.clear() + } + + // Validate the data format + if (!data || !data.nouns || !data.verbs || !data.version) { + throw new Error('Invalid restore data format') + } + + // Log additional data if present + if (data.nounTypes) { + console.log(`Found ${data.nounTypes.length} noun types in restore data`) + } + + if (data.verbTypes) { + console.log(`Found ${data.verbTypes.length} verb types in restore data`) + } + + if (data.hnswIndex) { + console.log('Found HNSW index data in backup') + } + + // Restore nouns + let nounsRestored = 0 + for (const noun of data.nouns) { + try { + // Check if the noun has a vector + if (!noun.vector || noun.vector.length === 0) { + // If no vector, create one using the embedding function + if ( + noun.metadata && + typeof noun.metadata === 'object' && + 'text' in noun.metadata + ) { + // If the metadata has a text field, use it for embedding + noun.vector = await this.embeddingFunction(noun.metadata.text) + } else { + // Otherwise, use the entire metadata for embedding + noun.vector = await this.embeddingFunction(noun.metadata) + } + } + + // Add the noun with its vector and metadata + await this.add(noun.vector, noun.metadata, { id: noun.id }) + nounsRestored++ + } catch (error) { + console.error(`Failed to restore noun ${noun.id}:`, error) + // Continue with other nouns + } + } + + // Restore verbs + let verbsRestored = 0 + for (const verb of data.verbs) { + try { + // Check if the verb has a vector + if (!verb.vector || verb.vector.length === 0) { + // If no vector, create one using the embedding function + if ( + verb.metadata && + typeof verb.metadata === 'object' && + 'text' in verb.metadata + ) { + // If the metadata has a text field, use it for embedding + verb.vector = await this.embeddingFunction(verb.metadata.text) + } else { + // Otherwise, use the entire metadata for embedding + verb.vector = await this.embeddingFunction(verb.metadata) + } + } + + // Add the verb + await this.addVerb(verb.sourceId, verb.targetId, verb.vector, { + id: verb.id, + type: verb.metadata?.verb || VerbType.RelatedTo, + metadata: verb.metadata + }) + verbsRestored++ + } catch (error) { + console.error(`Failed to restore verb ${verb.id}:`, error) + // Continue with other verbs + } + } + + // If HNSW index data is provided and we've restored nouns, reconstruct the index + if (data.hnswIndex && nounsRestored > 0) { + try { + console.log('Reconstructing HNSW index from backup data...') + + // Create a new index with the restored configuration + this.index = new HNSWIndex( + data.hnswIndex.config, + this.distanceFunction + ) + + // Re-add all nouns to the index + for (const noun of data.nouns) { + if (noun.vector && noun.vector.length > 0) { + this.index.addItem({ id: noun.id, vector: noun.vector }) + } + } + + console.log('HNSW index reconstruction complete') + } catch (error) { + console.error('Failed to reconstruct HNSW index:', error) + console.log('Continuing with standard restore process...') + } + } + + return { + nounsRestored, + verbsRestored + } + } catch (error) { + console.error('Failed to restore data:', error) + throw new Error(`Failed to restore data: ${error}`) + } + } + /** * Generate a random graph of data with typed nouns and verbs for testing and experimentation * @param options Configuration options for the random graph * @returns Object containing the IDs of the generated nouns and verbs */ - public async generateRandomGraph(options: { - nounCount?: number; // Number of nouns to generate (default: 10) - verbCount?: number; // Number of verbs to generate (default: 20) - nounTypes?: NounType[]; // Types of nouns to generate (default: all types) - verbTypes?: VerbType[]; // Types of verbs to generate (default: all types) - clearExisting?: boolean; // Whether to clear existing data before generating (default: false) - seed?: string; // Seed for random generation (default: random) - } = {}): Promise<{ - nounIds: string[]; - verbIds: string[]; + public async generateRandomGraph( + options: { + nounCount?: number // Number of nouns to generate (default: 10) + verbCount?: number // Number of verbs to generate (default: 20) + nounTypes?: NounType[] // Types of nouns to generate (default: all types) + verbTypes?: VerbType[] // Types of verbs to generate (default: all types) + clearExisting?: boolean // Whether to clear existing data before generating (default: false) + seed?: string // Seed for random generation (default: random) + } = {} + ): Promise<{ + nounIds: string[] + verbIds: string[] }> { await this.ensureInitialized() @@ -1491,7 +1897,8 @@ export class BrainyData { // Create metadata const metadata = { verb: verbType, - description: verbDescriptions[verbType] || `A random ${verbType} relationship`, + description: + verbDescriptions[verbType] || `A random ${verbType} relationship`, weight: Math.random(), confidence: Math.random(), randomAttributes: { @@ -1523,4 +1930,9 @@ export class BrainyData { } // Export distance functions for convenience -export { euclideanDistance, cosineDistance, manhattanDistance, dotProductDistance } from './utils/index.js' +export { + euclideanDistance, + cosineDistance, + manhattanDistance, + dotProductDistance +} from './utils/index.js' diff --git a/src/hnsw/hnswIndex.ts b/src/hnsw/hnswIndex.ts index ee4ff753..0d89879f 100644 --- a/src/hnsw/hnswIndex.ts +++ b/src/hnsw/hnswIndex.ts @@ -379,6 +379,34 @@ export class HNSWIndex { return this.distanceFunction } + /** + * Get the entry point ID + */ + public getEntryPointId(): string | null { + return this.entryPointId + } + + /** + * Get the maximum level + */ + public getMaxLevel(): number { + return this.maxLevel + } + + /** + * Get the dimension + */ + public getDimension(): number | null { + return this.dimension + } + + /** + * Get the configuration + */ + public getConfig(): HNSWConfig { + return { ...this.config } + } + /** * Search within a specific layer * Returns a map of noun IDs to distances, sorted by distance diff --git a/src/hnsw/hnswIndexOptimized.ts b/src/hnsw/hnswIndexOptimized.ts new file mode 100644 index 00000000..f1c240a6 --- /dev/null +++ b/src/hnsw/hnswIndexOptimized.ts @@ -0,0 +1,586 @@ +/** + * Optimized HNSW (Hierarchical Navigable Small World) Index implementation + * Extends the base HNSW implementation with support for large datasets + * Uses product quantization for dimensionality reduction and disk-based storage when needed + */ + +import { + DistanceFunction, + HNSWConfig, + HNSWNoun, + Vector, + VectorDocument +} from '../coreTypes.js' +import { HNSWIndex } from './hnswIndex.js' +import { StorageAdapter } from '../coreTypes.js' + +// Configuration for the optimized HNSW index +export interface HNSWOptimizedConfig extends HNSWConfig { + // Memory threshold in bytes - when exceeded, will use disk-based approach + memoryThreshold?: number + + // Product quantization settings + productQuantization?: { + // Whether to use product quantization + enabled: boolean + // Number of subvectors to split the vector into + numSubvectors?: number + // Number of centroids per subvector + numCentroids?: number + } + + // Whether to use disk-based storage for the index + useDiskBasedIndex?: boolean +} + +// Default configuration for the optimized HNSW index +const DEFAULT_OPTIMIZED_CONFIG: HNSWOptimizedConfig = { + M: 16, + efConstruction: 200, + efSearch: 50, + ml: 16, + memoryThreshold: 1024 * 1024 * 1024, // 1GB default threshold + productQuantization: { + enabled: false, + numSubvectors: 16, + numCentroids: 256 + }, + useDiskBasedIndex: false +} + +/** + * Product Quantization implementation + * Reduces vector dimensionality by splitting vectors into subvectors + * and quantizing each subvector to the nearest centroid + */ +class ProductQuantizer { + private numSubvectors: number + private numCentroids: number + private centroids: Vector[][] = [] + private subvectorSize: number = 0 + private initialized: boolean = false + private dimension: number = 0 + + constructor(numSubvectors: number = 16, numCentroids: number = 256) { + this.numSubvectors = numSubvectors + this.numCentroids = numCentroids + } + + /** + * Initialize the product quantizer with training data + * @param vectors Training vectors to use for learning centroids + */ + public train(vectors: Vector[]): void { + if (vectors.length === 0) { + throw new Error('Cannot train product quantizer with empty vector set') + } + + this.dimension = vectors[0].length + this.subvectorSize = Math.ceil(this.dimension / this.numSubvectors) + + // Initialize centroids for each subvector + for (let i = 0; i < this.numSubvectors; i++) { + // Extract subvectors from training data + const subvectors: Vector[] = vectors.map((vector) => { + const start = i * this.subvectorSize + const end = Math.min(start + this.subvectorSize, this.dimension) + return vector.slice(start, end) + }) + + // Initialize centroids for this subvector using k-means++ + this.centroids[i] = this.kMeansPlusPlus(subvectors, this.numCentroids) + } + + this.initialized = true + } + + /** + * Quantize a vector using product quantization + * @param vector Vector to quantize + * @returns Array of centroid indices, one for each subvector + */ + public quantize(vector: Vector): number[] { + if (!this.initialized) { + throw new Error('Product quantizer not initialized. Call train() first.') + } + + if (vector.length !== this.dimension) { + throw new Error( + `Vector dimension mismatch: expected ${this.dimension}, got ${vector.length}` + ) + } + + const codes: number[] = [] + + // Quantize each subvector + for (let i = 0; i < this.numSubvectors; i++) { + const start = i * this.subvectorSize + const end = Math.min(start + this.subvectorSize, this.dimension) + const subvector = vector.slice(start, end) + + // Find nearest centroid + let minDist = Number.MAX_VALUE + let nearestCentroidIndex = 0 + + for (let j = 0; j < this.centroids[i].length; j++) { + const centroid = this.centroids[i][j] + const dist = this.euclideanDistanceSquared(subvector, centroid) + + if (dist < minDist) { + minDist = dist + nearestCentroidIndex = j + } + } + + codes.push(nearestCentroidIndex) + } + + return codes + } + + /** + * Reconstruct a vector from its quantized representation + * @param codes Array of centroid indices + * @returns Reconstructed vector + */ + public reconstruct(codes: number[]): Vector { + if (!this.initialized) { + throw new Error('Product quantizer not initialized. Call train() first.') + } + + if (codes.length !== this.numSubvectors) { + throw new Error( + `Code length mismatch: expected ${this.numSubvectors}, got ${codes.length}` + ) + } + + const reconstructed: Vector = [] + + // Reconstruct each subvector + for (let i = 0; i < this.numSubvectors; i++) { + const centroidIndex = codes[i] + const centroid = this.centroids[i][centroidIndex] + + // Add centroid components to reconstructed vector + for (const component of centroid) { + reconstructed.push(component) + } + } + + // Trim to original dimension if needed + return reconstructed.slice(0, this.dimension) + } + + /** + * Compute squared Euclidean distance between two vectors + * @param a First vector + * @param b Second vector + * @returns Squared Euclidean distance + */ + private euclideanDistanceSquared(a: Vector, b: Vector): number { + let sum = 0 + const length = Math.min(a.length, b.length) + + for (let i = 0; i < length; i++) { + const diff = a[i] - b[i] + sum += diff * diff + } + + return sum + } + + /** + * Implement k-means++ algorithm to initialize centroids + * @param vectors Vectors to cluster + * @param k Number of clusters + * @returns Array of centroids + */ + private kMeansPlusPlus(vectors: Vector[], k: number): Vector[] { + if (vectors.length < k) { + // If we have fewer vectors than centroids, use the vectors as centroids + return [...vectors] + } + + const centroids: Vector[] = [] + + // Choose first centroid randomly + const firstIndex = Math.floor(Math.random() * vectors.length) + centroids.push([...vectors[firstIndex]]) + + // Choose remaining centroids + for (let i = 1; i < k; i++) { + // Compute distances to nearest centroid for each vector + const distances: number[] = vectors.map((vector) => { + let minDist = Number.MAX_VALUE + + for (const centroid of centroids) { + const dist = this.euclideanDistanceSquared(vector, centroid) + minDist = Math.min(minDist, dist) + } + + return minDist + }) + + // Compute sum of distances + const distSum = distances.reduce((sum, dist) => sum + dist, 0) + + // Choose next centroid with probability proportional to distance + let r = Math.random() * distSum + let nextIndex = 0 + + for (let j = 0; j < distances.length; j++) { + r -= distances[j] + if (r <= 0) { + nextIndex = j + break + } + } + + centroids.push([...vectors[nextIndex]]) + } + + return centroids + } + + /** + * Get the centroids for each subvector + * @returns Array of centroid arrays + */ + public getCentroids(): Vector[][] { + return this.centroids + } + + /** + * Set the centroids for each subvector + * @param centroids Array of centroid arrays + */ + public setCentroids(centroids: Vector[][]): void { + this.centroids = centroids + this.numSubvectors = centroids.length + this.numCentroids = centroids[0].length + this.initialized = true + } + + /** + * Get the dimension of the vectors + * @returns Dimension + */ + public getDimension(): number { + return this.dimension + } + + /** + * Set the dimension of the vectors + * @param dimension Dimension + */ + public setDimension(dimension: number): void { + this.dimension = dimension + this.subvectorSize = Math.ceil(dimension / this.numSubvectors) + } +} + +/** + * Optimized HNSW Index implementation + * Extends the base HNSW implementation with support for large datasets + * Uses product quantization for dimensionality reduction and disk-based storage when needed + */ +export class HNSWIndexOptimized extends HNSWIndex { + private optimizedConfig: HNSWOptimizedConfig + private productQuantizer: ProductQuantizer | null = null + private storage: StorageAdapter | null = null + private useDiskBasedIndex: boolean = false + private useProductQuantization: boolean = false + private quantizedVectors: Map = new Map() + private memoryUsage: number = 0 + private vectorCount: number = 0 + + constructor( + config: Partial = {}, + distanceFunction: DistanceFunction, + storage: StorageAdapter | null = null + ) { + // Initialize base HNSW index with standard config + super(config, distanceFunction) + + // Set optimized config + this.optimizedConfig = { ...DEFAULT_OPTIMIZED_CONFIG, ...config } + + // Set storage adapter + this.storage = storage + + // Initialize product quantizer if enabled + if (this.optimizedConfig.productQuantization?.enabled) { + this.useProductQuantization = true + this.productQuantizer = new ProductQuantizer( + this.optimizedConfig.productQuantization.numSubvectors, + this.optimizedConfig.productQuantization.numCentroids + ) + } + + // Set disk-based index flag + this.useDiskBasedIndex = this.optimizedConfig.useDiskBasedIndex || false + } + + /** + * Add a vector to the index + * Uses product quantization if enabled and memory threshold is exceeded + */ + public override addItem(item: VectorDocument): string { + // Check if item is defined + if (!item) { + throw new Error('Item is undefined or null') + } + + const { id, vector } = item + + // Check if vector is defined + if (!vector) { + throw new Error('Vector is undefined or null') + } + + // Estimate memory usage for this vector + const vectorMemory = vector.length * 8 // 8 bytes per number (Float64) + const connectionsMemory = this.optimizedConfig.M * this.optimizedConfig.ml * 16 // Estimate for connections + const totalMemory = vectorMemory + connectionsMemory + + // Update memory usage estimate + this.memoryUsage += totalMemory + this.vectorCount++ + + // Check if we should switch to product quantization + if ( + this.useProductQuantization && + this.memoryUsage > this.optimizedConfig.memoryThreshold! && + this.productQuantizer && + !this.productQuantizer.getDimension() + ) { + // Initialize product quantizer with existing vectors + this.initializeProductQuantizer() + } + + // If product quantization is active, quantize the vector + if ( + this.useProductQuantization && + this.productQuantizer && + this.productQuantizer.getDimension() > 0 + ) { + // Quantize the vector + const codes = this.productQuantizer.quantize(vector) + + // Store the quantized vector + this.quantizedVectors.set(id, codes) + + // Reconstruct the vector for indexing + const reconstructedVector = this.productQuantizer.reconstruct(codes) + + // Add the reconstructed vector to the index + return super.addItem({ id, vector: reconstructedVector }) + } + + // If disk-based index is active and storage is available, store the vector + if (this.useDiskBasedIndex && this.storage) { + // Create a noun object + const noun: HNSWNoun = { + id, + vector, + connections: new Map() + } + + // Store the noun + this.storage.saveNoun(noun).catch((error) => { + console.error(`Failed to save noun ${id} to storage:`, error) + }) + } + + // Add the vector to the in-memory index + return super.addItem(item) + } + + /** + * Search for nearest neighbors + * Uses product quantization if enabled + */ + public override search( + queryVector: Vector, + k: number = 10 + ): Array<[string, number]> { + // Check if query vector is defined + if (!queryVector) { + throw new Error('Query vector is undefined or null') + } + + // If product quantization is active, quantize the query vector + if ( + this.useProductQuantization && + this.productQuantizer && + this.productQuantizer.getDimension() > 0 + ) { + // Quantize the query vector + const codes = this.productQuantizer.quantize(queryVector) + + // Reconstruct the query vector + const reconstructedVector = this.productQuantizer.reconstruct(codes) + + // Search with the reconstructed vector + return super.search(reconstructedVector, k) + } + + // Otherwise, use the standard search + return super.search(queryVector, k) + } + + /** + * Remove an item from the index + */ + public override removeItem(id: string): boolean { + // If product quantization is active, remove the quantized vector + if (this.useProductQuantization) { + this.quantizedVectors.delete(id) + } + + // If disk-based index is active and storage is available, remove the vector from storage + if (this.useDiskBasedIndex && this.storage) { + this.storage.deleteNoun(id).catch((error) => { + console.error(`Failed to delete noun ${id} from storage:`, error) + }) + } + + // Update memory usage estimate + if (this.vectorCount > 0) { + this.memoryUsage = Math.max( + 0, + this.memoryUsage - this.memoryUsage / this.vectorCount + ) + this.vectorCount-- + } + + // Remove the item from the in-memory index + return super.removeItem(id) + } + + /** + * Clear the index + */ + public override clear(): void { + // Clear product quantization data + if (this.useProductQuantization) { + this.quantizedVectors.clear() + this.productQuantizer = new ProductQuantizer( + this.optimizedConfig.productQuantization!.numSubvectors, + this.optimizedConfig.productQuantization!.numCentroids + ) + } + + // Reset memory usage + this.memoryUsage = 0 + this.vectorCount = 0 + + // Clear the in-memory index + super.clear() + } + + /** + * Initialize product quantizer with existing vectors + */ + private initializeProductQuantizer(): void { + if (!this.productQuantizer) { + return + } + + // Get all vectors from the index + const nouns = super.getNouns() + const vectors: Vector[] = [] + + // Extract vectors + for (const [_, noun] of nouns) { + vectors.push(noun.vector) + } + + // Train the product quantizer + if (vectors.length > 0) { + this.productQuantizer.train(vectors) + + // Quantize all existing vectors + for (const [id, noun] of nouns) { + const codes = this.productQuantizer.quantize(noun.vector) + this.quantizedVectors.set(id, codes) + } + + console.log( + `Initialized product quantizer with ${vectors.length} vectors` + ) + } + } + + /** + * Get the product quantizer + * @returns Product quantizer or null if not enabled + */ + public getProductQuantizer(): ProductQuantizer | null { + return this.productQuantizer + } + + /** + * Get the optimized configuration + * @returns Optimized configuration + */ + public getOptimizedConfig(): HNSWOptimizedConfig { + return { ...this.optimizedConfig } + } + + /** + * Get the estimated memory usage + * @returns Estimated memory usage in bytes + */ + public getMemoryUsage(): number { + return this.memoryUsage + } + + /** + * Set the storage adapter + * @param storage Storage adapter + */ + public setStorage(storage: StorageAdapter): void { + this.storage = storage + } + + /** + * Get the storage adapter + * @returns Storage adapter or null if not set + */ + public getStorage(): StorageAdapter | null { + return this.storage + } + + /** + * Set whether to use disk-based index + * @param useDiskBasedIndex Whether to use disk-based index + */ + public setUseDiskBasedIndex(useDiskBasedIndex: boolean): void { + this.useDiskBasedIndex = useDiskBasedIndex + } + + /** + * Get whether disk-based index is used + * @returns Whether disk-based index is used + */ + public getUseDiskBasedIndex(): boolean { + return this.useDiskBasedIndex + } + + /** + * Set whether to use product quantization + * @param useProductQuantization Whether to use product quantization + */ + public setUseProductQuantization(useProductQuantization: boolean): void { + this.useProductQuantization = useProductQuantization + } + + /** + * Get whether product quantization is used + * @returns Whether product quantization is used + */ + public getUseProductQuantization(): boolean { + return this.useProductQuantization + } +}