soulcraft/brainy
2
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions
brainy/src/brainyData.ts
dpsifr bc480770b5 **feat(core, statistics, storage): enhance service-level statistics tracking and result filtering** 
- **Core**: Introduced a `getCurrentAugmentation` method for detecting active augmentation names. Updated metadata handling to include `createdBy`, `createdAt`, and `updatedAt` attributes for improved tracking.
- **Storage**: Added support for service-based statistics tracking with new methods such as `incrementStatistic`, `decrementStatistic`, and `updateHnswIndexSize`. Implemented persistence for statistics in storage adapters.
- **Statistics**: Enhanced `getStatistics` functionality to provide service-specific breakdowns and support filtering by services. Improved noun, verb, and metadata tracking mechanisms.
- **Search**: Added `service` option to filter results during searches for nouns, verbs, and metadata, ensuring accurate service-based query results.
- **Refactor**: Simplified search logic by integrating HNSW index filtering for better performance when retrieving service-specific results.
- **Tests**: Added comprehensive test coverage for service-level statistics and filtering by service.

**Purpose**: Improve service-level data tracking and analytics while enhancing functionality for filtering and maintaining metadata accuracy to support detailed insights for diverse use cases.
2025-07-24 12:07:47 -07:00

3060 lines
112 KiB
TypeScript
Raw Blame History

/**
* BrainyData
* Main class that provides the vector database functionality
*/
import {v4 as uuidv4} from 'uuid'
import {HNSWIndex} from './hnsw/hnswIndex.js'
import {
HNSWIndexOptimized,
HNSWOptimizedConfig
} from './hnsw/hnswIndexOptimized.js'
import {createStorage} from './storage/storageFactory.js'
import {
DistanceFunction,
GraphVerb,
EmbeddingFunction,
HNSWConfig,
HNSWNoun,
SearchResult,
StorageAdapter,
Vector,
VectorDocument
} from './coreTypes.js'
import {
cosineDistance,
defaultEmbeddingFunction,
defaultBatchEmbeddingFunction,
getDefaultEmbeddingFunction,
getDefaultBatchEmbeddingFunction,
euclideanDistance,
cleanupWorkerPools
} from './utils/index.js'
import {NounType, VerbType, GraphNoun} from './types/graphTypes.js'
import {
ServerSearchConduitAugmentation,
createServerSearchAugmentations
} from './augmentations/serverSearchAugmentations.js'
import {WebSocketConnection, AugmentationType, IAugmentation} from './types/augmentations.js'
import {BrainyDataInterface} from './types/brainyDataInterface.js'
import {augmentationPipeline} from './augmentationPipeline.js'
export interface BrainyDataConfig {
/**
* Vector dimensions (required if not using an embedding function that auto-detects dimensions)
*/
dimensions?: number
/**
* HNSW index configuration
*/
hnsw?: Partial<HNSWConfig>
/**
* Optimized HNSW index configuration
* If provided, will use the optimized HNSW index instead of the standard one
*/
hnswOptimized?: Partial<HNSWOptimizedConfig>
/**
* Distance function to use for similarity calculations
*/
distanceFunction?: DistanceFunction
/**
* Custom storage adapter (if not provided, will use OPFS or memory storage)
*/
storageAdapter?: StorageAdapter
/**
* Storage configuration options
* These will be passed to createStorage if storageAdapter is not provided
*/
storage?: {
requestPersistentStorage?: boolean
r2Storage?: {
bucketName?: string
accountId?: string
accessKeyId?: string
secretAccessKey?: string
}
s3Storage?: {
bucketName?: string
accessKeyId?: string
secretAccessKey?: string
region?: string
}
gcsStorage?: {
bucketName?: string
accessKeyId?: string
secretAccessKey?: string
endpoint?: string
}
customS3Storage?: {
bucketName?: string
accessKeyId?: string
secretAccessKey?: string
endpoint?: string
region?: string
}
forceFileSystemStorage?: boolean
forceMemoryStorage?: boolean
}
/**
* Embedding function to convert data to vectors
*/
embeddingFunction?: EmbeddingFunction
/**
* Set the database to read-only mode
* When true, all write operations will throw an error
*/
readOnly?: boolean
/**
* Remote server configuration for search operations
*/
remoteServer?: {
/**
* WebSocket URL of the remote Brainy server
*/
url: string
/**
* WebSocket protocols to use for the connection
*/
protocols?: string | string[]
/**
* Whether to automatically connect to the remote server on initialization
*/
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> {
private index: HNSWIndex | HNSWIndexOptimized
private storage: StorageAdapter | null = null
private isInitialized = false
private isInitializing = false
private embeddingFunction: EmbeddingFunction
private distanceFunction: DistanceFunction
private requestPersistentStorage: boolean
private readOnly: boolean
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
private serverSearchConduit: ServerSearchConduitAugmentation | null = null
private serverConnection: WebSocketConnection | null = null
/**
* Get the vector dimensions
*/
public get dimensions(): number {
return this._dimensions
}
/**
* Get the maximum connections parameter from HNSW configuration
*/
public get maxConnections(): number {
const config = this.index.getConfig()
return config.M || 16
}
/**
* Get the efConstruction parameter from HNSW configuration
*/
public get efConstruction(): number {
const config = this.index.getConfig()
return config.efConstruction || 200
}
/**
* Create a new vector database
*/
constructor(config: BrainyDataConfig = {}) {
// Validate dimensions
if (config.dimensions !== undefined && config.dimensions <= 0) {
throw new Error('Dimensions must be a positive number')
}
// Set dimensions (default to 512 for embedding functions, or require explicit config)
this._dimensions = config.dimensions || 512
// 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
// 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 =
config.storage?.requestPersistentStorage || false
// Set read-only flag
this.readOnly = config.readOnly || false
// Store storage configuration for later use in init()
this.storageConfig = config.storage || {}
// Store remote server configuration if provided
if (config.remoteServer) {
this.remoteServerConfig = config.remoteServer
}
}
/**
* Check if the database is in read-only mode and throw an error if it is
* @throws Error if the database is in read-only mode
*/
private checkReadOnly(): void {
if (this.readOnly) {
throw new Error(
'Cannot perform write operation: database is in read-only mode'
)
}
}
/**
* Get the current augmentation name if available
* This is used to auto-detect the service performing data operations
* @returns The name of the current augmentation or 'default' if none is detected
*/
private getCurrentAugmentation(): string {
try {
// Get all registered augmentations
const augmentationTypes = augmentationPipeline.getAvailableAugmentationTypes()
// Check each type of augmentation
for (const type of augmentationTypes) {
const augmentations = augmentationPipeline.getAugmentationsByType(type)
// Find the first enabled augmentation
for (const augmentation of augmentations) {
if (augmentation.enabled) {
return augmentation.name
}
}
}
return 'default'
} catch (error) {
// If there's any error in detection, return default
console.warn('Failed to detect current augmentation:', error)
return 'default'
}
}
/**
* Initialize the database
* Loads existing data from storage if available
*/
public async init(): Promise<void> {
if (this.isInitialized) {
return
}
// Prevent recursive initialization
if (this.isInitializing) {
return
}
this.isInitializing = true
try {
// Pre-load the embedding model early to ensure it's always available
// This helps prevent issues with the Universal Sentence Encoder not being loaded
try {
// Pre-loading Universal Sentence Encoder model
// Call embedding function directly to avoid circular dependency with embed()
await this.embeddingFunction('')
// Universal Sentence Encoder model loaded successfully
} catch (embedError) {
console.warn(
'Failed to pre-load Universal Sentence Encoder:',
embedError
)
// Try again with a retry mechanism
// Retrying Universal Sentence Encoder initialization
try {
// Wait a moment before retrying
await new Promise((resolve) => setTimeout(resolve, 1000))
// Try again with a different approach - use the non-threaded version
// This is a fallback in case the threaded version fails
const {createTensorFlowEmbeddingFunction} = await import(
'./utils/embedding.js'
)
const fallbackEmbeddingFunction = createTensorFlowEmbeddingFunction()
// Test the fallback embedding function
await fallbackEmbeddingFunction('')
// If successful, replace the embedding function
console.log(
'Successfully loaded Universal Sentence Encoder with fallback method'
)
this.embeddingFunction = fallbackEmbeddingFunction
} catch (retryError) {
console.error(
'All attempts to load Universal Sentence Encoder failed:',
retryError
)
// Continue initialization even if embedding model fails to load
// The application will need to handle missing embedding functionality
}
}
// Initialize storage if not provided in constructor
if (!this.storage) {
// Combine storage config with requestPersistentStorage for backward compatibility
let storageOptions = {
...this.storageConfig,
requestPersistentStorage: this.requestPersistentStorage
}
// Ensure s3Storage has all required fields if it's provided
if (storageOptions.s3Storage) {
// Only include s3Storage if all required fields are present
if (storageOptions.s3Storage.bucketName &&
storageOptions.s3Storage.accessKeyId &&
storageOptions.s3Storage.secretAccessKey) {
// All required fields are present, keep s3Storage as is
} else {
// Missing required fields, remove s3Storage to avoid type errors
const {s3Storage, ...rest} = storageOptions
storageOptions = rest
console.warn('Ignoring s3Storage configuration due to missing required fields')
}
}
// Use type assertion to tell TypeScript that storageOptions conforms to StorageOptions
this.storage = await createStorage(storageOptions as any)
}
// 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()
// Clear the index and add all nouns
this.index.clear()
for (const noun of nouns) {
// Check if the vector dimensions match the expected dimensions
if (noun.vector.length !== this._dimensions) {
console.warn(
`Skipping noun ${noun.id} due to dimension mismatch: expected ${this._dimensions}, got ${noun.vector.length}`
)
// Optionally, you could delete the mismatched noun from storage
// await this.storage!.deleteNoun(noun.id)
continue
}
// Add to index
await this.index.addItem({
id: noun.id,
vector: noun.vector
})
}
// Connect to remote server if configured with autoConnect
if (this.remoteServerConfig && this.remoteServerConfig.autoConnect) {
try {
await this.connectToRemoteServer(
this.remoteServerConfig.url,
this.remoteServerConfig.protocols
)
} catch (remoteError) {
console.warn('Failed to auto-connect to remote server:', remoteError)
// Continue initialization even if remote connection fails
}
}
this.isInitialized = true
this.isInitializing = false
} catch (error) {
console.error('Failed to initialize BrainyData:', error)
this.isInitializing = false
throw new Error(`Failed to initialize BrainyData: ${error}`)
}
}
/**
* Connect to a remote Brainy server for search operations
* @param serverUrl WebSocket URL of the remote Brainy server
* @param protocols Optional WebSocket protocols to use
* @returns The connection object
*/
public async connectToRemoteServer(
serverUrl: string,
protocols?: string | string[]
): Promise<WebSocketConnection> {
await this.ensureInitialized()
try {
// Create server search augmentations
const {conduit, connection} = await createServerSearchAugmentations(
serverUrl,
{
protocols,
localDb: this
}
)
// Store the conduit and connection
this.serverSearchConduit = conduit
this.serverConnection = connection
return connection
} catch (error) {
console.error('Failed to connect to remote server:', error)
throw new Error(`Failed to connect to remote server: ${error}`)
}
}
/**
* Add a vector or data to the database
* If the input is not a vector, it will be converted using the embedding function
* @param vectorOrData Vector or data to add
* @param metadata Optional metadata to associate with the vector
* @param options Additional options
* @returns The ID of the added vector
*/
public async add(
vectorOrData: Vector | any,
metadata?: T,
options: {
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
service?: string // The service that is inserting the data
} = {}
): Promise<string> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
let vector: Vector
// Check if input is already a vector
if (
Array.isArray(vectorOrData) &&
vectorOrData.every((item) => typeof item === 'number') &&
!options.forceEmbed
) {
// Input is already a vector
vector = vectorOrData
} else {
// Input needs to be vectorized
try {
vector = await this.embeddingFunction(vectorOrData)
} catch (embedError) {
throw new Error(`Failed to vectorize data: ${embedError}`)
}
}
// Check if vector is defined
if (!vector) {
throw new Error('Vector is undefined or null')
}
// Validate vector dimensions
if (vector.length !== this._dimensions) {
throw new Error(`Vector dimension mismatch: expected ${this._dimensions}, got ${vector.length}`)
}
// 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
await this.index.addItem({id, vector})
// Get the noun from the index
const noun = this.index.getNouns().get(id)
if (!noun) {
throw new Error(`Failed to retrieve newly created noun with ID ${id}`)
}
// Save noun to storage
await this.storage!.saveNoun(noun)
// Track noun statistics
const service = options.service || this.getCurrentAugmentation()
await this.storage!.incrementStatistic('noun', service)
// Save metadata if provided
if (metadata !== undefined) {
// Validate noun type if metadata is for a GraphNoun
if (metadata && typeof metadata === 'object' && 'noun' in metadata) {
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.`
)
// Set a default noun type
;(metadata as unknown as GraphNoun).noun = NounType.Concept
}
// Ensure createdBy field is populated for GraphNoun
const service = options.service || this.getCurrentAugmentation()
const graphNoun = metadata as unknown as GraphNoun
// Only set createdBy if it doesn't exist or is being explicitly updated
if (!graphNoun.createdBy || options.service) {
graphNoun.createdBy = {
augmentation: service,
version: '1.0' // TODO: Get actual version from augmentation
}
}
// Update timestamps
const now = new Date()
const timestamp = {
seconds: Math.floor(now.getTime() / 1000),
nanoseconds: (now.getTime() % 1000) * 1000000
}
// Set createdAt if it doesn't exist
if (!graphNoun.createdAt) {
graphNoun.createdAt = timestamp
}
// Always update updatedAt
graphNoun.updatedAt = timestamp
}
// Ensure metadata has the correct id field
let metadataToSave = metadata
if (metadata && typeof metadata === 'object') {
metadataToSave = {...metadata, id}
}
await this.storage!.saveMetadata(id, metadataToSave)
// Track metadata statistics
const metadataService = options.service || this.getCurrentAugmentation()
await this.storage!.incrementStatistic('metadata', metadataService)
}
// Update HNSW index size (excluding verbs)
await this.storage!.updateHnswIndexSize(await this.getNounCount())
// If addToRemote is true and we're connected to a remote server, add to remote as well
if (options.addToRemote && this.isConnectedToRemoteServer()) {
try {
await this.addToRemote(id, vector, metadata)
} catch (remoteError) {
console.warn(
`Failed to add to remote server: ${remoteError}. Continuing with local add.`
)
}
}
return id
} catch (error) {
console.error('Failed to add vector:', error)
throw new Error(`Failed to add vector: ${error}`)
}
}
/**
* Add a text item to the database with automatic embedding
* This is a convenience method for adding text data with metadata
* @param text Text data to add
* @param metadata Metadata to associate with the text
* @param options Additional options
* @returns The ID of the added item
*/
public async addItem(
text: string,
metadata?: T,
options: {
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<string> {
// Use the existing add method with forceEmbed to ensure text is embedded
return this.add(text, metadata, {...options, forceEmbed: true})
}
/**
* Add data to both local and remote Brainy instances
* @param vectorOrData Vector or data to add
* @param metadata Optional metadata to associate with the vector
* @param options Additional options
* @returns The ID of the added vector
*/
public async addToBoth(
vectorOrData: Vector | any,
metadata?: T,
options: {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
} = {}
): Promise<string> {
// Check if connected to a remote server
if (!this.isConnectedToRemoteServer()) {
throw new Error(
'Not connected to a remote server. Call connectToRemoteServer() first.'
)
}
// Add to local with addToRemote option
return this.add(vectorOrData, metadata, {...options, addToRemote: true})
}
/**
* Add a vector to the remote server
* @param id ID of the vector to add
* @param vector Vector to add
* @param metadata Optional metadata to associate with the vector
* @returns True if successful, false otherwise
* @private
*/
private async addToRemote(
id: string,
vector: Vector,
metadata?: T
): Promise<boolean> {
if (!this.isConnectedToRemoteServer()) {
return false
}
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,
vector,
metadata
)
if (!addResult.success) {
throw new Error(`Remote add failed: ${addResult.error}`)
}
return true
} catch (error) {
console.error('Failed to add to remote server:', error)
throw new Error(`Failed to add to remote server: ${error}`)
}
}
/**
* Add multiple vectors or data items to the database
* @param items Array of items to add
* @param options Additional options
* @returns Array of IDs for the added items
*/
public async addBatch(
items: Array<{
vectorOrData: Vector | any
metadata?: T
}>,
options: {
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
concurrency?: number // Maximum number of concurrent operations (default: 4)
batchSize?: number // Maximum number of items to process in a single batch (default: 50)
} = {}
): Promise<string[]> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
// Default concurrency to 4 if not specified
const concurrency = options.concurrency || 4
// Default batch size to 50 if not specified
const batchSize = options.batchSize || 50
try {
// Process items in batches to control concurrency and memory usage
const ids: string[] = []
const itemsToProcess = [...items] // Create a copy to avoid modifying the original array
while (itemsToProcess.length > 0) {
// Take up to 'batchSize' items to process in a batch
const batch = itemsToProcess.splice(0, batchSize)
// Separate items that are already vectors from those that need embedding
const vectorItems: Array<{
vectorOrData: Vector
metadata?: T
index: number
}> = []
const textItems: Array<{
text: string
metadata?: T
index: number
}> = []
// Categorize items
batch.forEach((item, index) => {
if (
Array.isArray(item.vectorOrData) &&
item.vectorOrData.every((val) => typeof val === 'number') &&
!options.forceEmbed
) {
// Item is already a vector
vectorItems.push({
vectorOrData: item.vectorOrData,
metadata: item.metadata,
index
})
} else if (typeof item.vectorOrData === 'string') {
// Item is text that needs embedding
textItems.push({
text: item.vectorOrData,
metadata: item.metadata,
index
})
} else {
// For now, treat other types as text
// In a more complete implementation, we might handle other types differently
const textRepresentation = String(item.vectorOrData)
textItems.push({
text: textRepresentation,
metadata: item.metadata,
index
})
}
})
// Process vector items (already embedded)
const vectorPromises = vectorItems.map((item) =>
this.add(item.vectorOrData, item.metadata, options)
)
// Process text items in a single batch embedding operation
let textPromises: Promise<string>[] = []
if (textItems.length > 0) {
// Extract just the text for batch embedding
const texts = textItems.map((item) => item.text)
// Perform batch embedding
const embeddings = await defaultBatchEmbeddingFunction(texts)
// Add each item with its embedding
textPromises = textItems.map((item, i) =>
this.add(embeddings[i], item.metadata, {
...options,
forceEmbed: false
})
)
}
// Combine all promises
const batchResults = await Promise.all([
...vectorPromises,
...textPromises
])
// Add the results to our ids array
ids.push(...batchResults)
}
return ids
} catch (error) {
console.error('Failed to add batch of items:', error)
throw new Error(`Failed to add batch of items: ${error}`)
}
}
/**
* Add multiple vectors or data items to both local and remote databases
* @param items Array of items to add
* @param options Additional options
* @returns Array of IDs for the added items
*/
public async addBatchToBoth(
items: Array<{
vectorOrData: Vector | any
metadata?: T
}>,
options: {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
concurrency?: number // Maximum number of concurrent operations (default: 4)
} = {}
): Promise<string[]> {
// Check if connected to a remote server
if (!this.isConnectedToRemoteServer()) {
throw new Error(
'Not connected to a remote server. Call connectToRemoteServer() first.'
)
}
// Add to local with addToRemote option
return this.addBatch(items, {...options, addToRemote: true})
}
/**
* Filter search results by service
* @param results Search results to filter
* @param service Service to filter by
* @returns Filtered search results
* @private
*/
private filterResultsByService<R extends SearchResult<T>>(
results: R[],
service?: string
): R[] {
if (!service) return results
return results.filter(result => {
if (!result.metadata || typeof result.metadata !== 'object') return false
if (!('createdBy' in result.metadata)) return false
const createdBy = result.metadata.createdBy as any
if (!createdBy) return false
return createdBy.augmentation === service
})
}
/**
* Search for similar vectors within specific noun types
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param nounTypes Array of noun types to search within, or null to search all
* @param options Additional options
* @returns Array of search results
*/
public async searchByNounTypes(
queryVectorOrData: Vector | any,
k: number = 10,
nounTypes: string[] | null = null,
options: {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
service?: string // Filter results by the service that created the data
} = {}
): Promise<SearchResult<T>[]> {
// Helper function to filter results by service
const filterByService = (metadata: any): boolean => {
if (!options.service) return true // No filter, include all
// Check if metadata has createdBy field with matching service
if (!metadata || typeof metadata !== 'object') return false
if (!('createdBy' in metadata)) return false
const createdBy = metadata.createdBy as any
if (!createdBy) return false
return createdBy.augmentation === options.service
}
if (!this.isInitialized) {
throw new Error('BrainyData must be initialized before searching. Call init() first.')
}
try {
let queryVector: Vector
// Check if input is already a vector
if (
Array.isArray(queryVectorOrData) &&
queryVectorOrData.every((item) => typeof item === 'number') &&
!options.forceEmbed
) {
// Input is already a vector
queryVector = queryVectorOrData
} else {
// Input needs to be vectorized
try {
queryVector = await this.embeddingFunction(queryVectorOrData)
} catch (embedError) {
throw new Error(`Failed to vectorize query data: ${embedError}`)
}
}
// Check if query vector is defined
if (!queryVector) {
throw new Error('Query vector is undefined or null')
}
// If no noun types specified, search all nouns
if (!nounTypes || nounTypes.length === 0) {
// Search in the index
const results = await this.index.search(queryVector, k)
// Get metadata for each result
const searchResults: SearchResult<T>[] = []
for (const [id, score] of results) {
const noun = this.index.getNouns().get(id)
if (!noun) {
continue
}
let metadata = await this.storage!.getMetadata(id)
// Initialize metadata to an empty object if it's null
if (metadata === null) {
metadata = {} as T
}
// Ensure metadata has the id field
if (metadata && typeof metadata === 'object') {
metadata = { ...metadata, id } as T
}
searchResults.push({
id,
score,
vector: noun.vector,
metadata: metadata as T
})
}
// Filter results by service if specified
return this.filterResultsByService(searchResults, options.service)
} else {
// Get nouns for each noun type in parallel
const nounPromises = nounTypes.map((nounType) =>
this.storage!.getNounsByNounType(nounType)
)
const nounArrays = await Promise.all(nounPromises)
// Combine all nouns
const nouns: HNSWNoun[] = []
for (const nounArray of nounArrays) {
nouns.push(...nounArray)
}
// Calculate distances for each noun
const results: Array<[string, number]> = []
for (const noun of nouns) {
const distance = this.index.getDistanceFunction()(
queryVector,
noun.vector
)
results.push([noun.id, distance])
}
// Sort by distance (ascending)
results.sort((a, b) => a[1] - b[1])
// Take top k results
const topResults = results.slice(0, k)
// Get metadata for each result
const searchResults: SearchResult<T>[] = []
for (const [id, score] of topResults) {
const noun = nouns.find((n) => n.id === id)
if (!noun) {
continue
}
let metadata = await this.storage!.getMetadata(id)
// Initialize metadata to an empty object if it's null
if (metadata === null) {
metadata = {} as T
}
// Ensure metadata has the id field
if (metadata && typeof metadata === 'object') {
metadata = { ...metadata, id } as T
}
searchResults.push({
id,
score,
vector: noun.vector,
metadata: metadata as T
})
}
// Filter results by service if specified
return this.filterResultsByService(searchResults, options.service)
}
} catch (error) {
console.error('Failed to search vectors by noun types:', error)
throw new Error(`Failed to search vectors by noun types: ${error}`)
}
}
/**
* Search for similar vectors
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async search(
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
searchMode?: 'local' | 'remote' | 'combined' // Where to search: local, remote, or both
searchVerbs?: boolean // Whether to search for verbs directly instead of nouns
verbTypes?: string[] // Optional array of verb types to search within or filter by
searchConnectedNouns?: boolean // Whether to search for nouns connected by verbs
verbDirection?: 'outgoing' | 'incoming' | 'both' // Direction of verbs to consider when searching connected nouns
service?: string // Filter results by the service that created the data
} = {}
): Promise<SearchResult<T>[]> {
if (!this.isInitialized) {
throw new Error('BrainyData must be initialized before searching. Call init() first.')
}
// If searching for verbs directly
if (options.searchVerbs) {
const verbResults = await this.searchVerbs(queryVectorOrData, k, {
forceEmbed: options.forceEmbed,
verbTypes: options.verbTypes
})
// Convert verb results to SearchResult format
return verbResults.map((verb) => ({
id: verb.id,
score: verb.similarity,
vector: verb.embedding || [],
metadata: {
verb: verb.verb,
source: verb.source,
target: verb.target,
...verb.data
} as unknown as T
}))
}
// If searching for nouns connected by verbs
if (options.searchConnectedNouns) {
return this.searchNounsByVerbs(queryVectorOrData, k, {
forceEmbed: options.forceEmbed,
verbTypes: options.verbTypes,
direction: options.verbDirection
})
}
// If a specific search mode is specified, use the appropriate search method
if (options.searchMode === 'local') {
return this.searchLocal(queryVectorOrData, k, options)
} else if (options.searchMode === 'remote') {
return this.searchRemote(queryVectorOrData, k, options)
} else if (options.searchMode === 'combined') {
return this.searchCombined(queryVectorOrData, k, options)
}
// Default behavior (backward compatible): search locally
return this.searchLocal(queryVectorOrData, k, options)
}
/**
* Search the local database for similar vectors
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async searchLocal(
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
service?: string // Filter results by the service that created the data
} = {}
): Promise<SearchResult<T>[]> {
if (!this.isInitialized) {
throw new Error('BrainyData must be initialized before searching. Call init() first.')
}
// If input is a string and not a vector, automatically vectorize it
let queryToUse = queryVectorOrData
if (typeof queryVectorOrData === 'string' && !options.forceEmbed) {
queryToUse = await this.embed(queryVectorOrData)
options.forceEmbed = false // Already embedded, don't force again
}
// 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,
service: options.service
}
)
} else {
// Otherwise, search all GraphNouns
searchResults = await this.searchByNounTypes(queryToUse, k, null, {
forceEmbed: options.forceEmbed,
service: options.service
})
}
// If includeVerbs is true, retrieve associated GraphVerbs for each result
if (options.includeVerbs && this.storage) {
for (const result of searchResults) {
try {
// Get outgoing verbs for this noun
const outgoingVerbs = await this.storage.getVerbsBySource(result.id)
// Get incoming verbs for this noun
const incomingVerbs = await this.storage.getVerbsByTarget(result.id)
// Combine all verbs
const allVerbs = [...outgoingVerbs, ...incomingVerbs]
// Add verbs to the result metadata
if (!result.metadata) {
result.metadata = {} as T
}
// Add the verbs to the metadata
;(result.metadata as Record<string, any>).associatedVerbs = allVerbs
} catch (error) {
console.warn(`Failed to retrieve verbs for noun ${result.id}:`, error)
}
}
}
return searchResults
}
/**
* Find entities similar to a given entity ID
* @param id ID of the entity to find similar entities for
* @param options Additional options
* @returns Array of search results with similarity scores
*/
public async findSimilar(
id: string,
options: {
limit?: number // Number of results to return
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<SearchResult<T>[]> {
await this.ensureInitialized()
// Get the entity by ID
const entity = await this.get(id)
if (!entity) {
throw new Error(`Entity with ID ${id} not found`)
}
// Use the entity's vector to search for similar entities
const k = (options.limit || 10) + 1 // Add 1 to account for the original entity
const searchResults = await this.search(entity.vector, k, {
forceEmbed: false,
nounTypes: options.nounTypes,
includeVerbs: options.includeVerbs,
searchMode: options.searchMode
})
// Filter out the original entity and limit to the requested number
return searchResults
.filter((result) => result.id !== id)
.slice(0, options.limit || 10)
}
/**
* Get a vector by ID
*/
public async get(id: string): Promise<VectorDocument<T> | null> {
await this.ensureInitialized()
try {
// Get noun from index
const noun = this.index.getNouns().get(id)
if (!noun) {
return null
}
// Get metadata
const metadata = await this.storage!.getMetadata(id)
return {
id,
vector: noun.vector,
metadata: metadata as T | undefined
}
} catch (error) {
console.error(`Failed to get vector ${id}:`, error)
throw new Error(`Failed to get vector ${id}: ${error}`)
}
}
/**
* Get all nouns in the database
* @returns Array of vector documents
*/
public async getAllNouns(): Promise<VectorDocument<T>[]> {
await this.ensureInitialized()
try {
const nouns = this.index.getNouns()
const result: VectorDocument<T>[] = []
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
* @param id The ID of the vector to delete
* @param options Additional options
* @returns Promise that resolves to true if the vector was deleted, false otherwise
*/
public async delete(
id: string,
options: {
service?: string // The service that is deleting the data
} = {}
): Promise<boolean> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
// Remove from index
const removed = this.index.removeItem(id)
if (!removed) {
return false
}
// Remove from storage
await this.storage!.deleteNoun(id)
// Track deletion statistics
const service = options.service || 'default'
await this.storage!.decrementStatistic('noun', service)
// Try to remove metadata (ignore errors)
try {
await this.storage!.saveMetadata(id, null)
await this.storage!.decrementStatistic('metadata', service)
} catch (error) {
// Ignore
}
return true
} catch (error) {
console.error(`Failed to delete vector ${id}:`, error)
throw new Error(`Failed to delete vector ${id}: ${error}`)
}
}
/**
* Update metadata for a vector
* @param id The ID of the vector to update metadata for
* @param metadata The new metadata
* @param options Additional options
* @returns Promise that resolves to true if the metadata was updated, false otherwise
*/
public async updateMetadata(
id: string,
metadata: T,
options: {
service?: string // The service that is updating the data
} = {}
): Promise<boolean> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
// Check if a vector exists
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 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.`
)
// Set a default noun type
;(metadata as unknown as GraphNoun).noun = NounType.Concept
}
// Get the service that's updating the metadata
const service = options.service || this.getCurrentAugmentation()
const graphNoun = metadata as unknown as GraphNoun
// Preserve existing createdBy and createdAt if they exist
const existingMetadata = await this.storage!.getMetadata(id) as any
if (existingMetadata &&
typeof existingMetadata === 'object' &&
'createdBy' in existingMetadata) {
// Preserve the original creator information
graphNoun.createdBy = existingMetadata.createdBy
// Also preserve creation timestamp if it exists
if ('createdAt' in existingMetadata) {
graphNoun.createdAt = existingMetadata.createdAt
}
} else if (!graphNoun.createdBy) {
// If no existing createdBy and none in the update, set it
graphNoun.createdBy = {
augmentation: service,
version: '1.0' // TODO: Get actual version from augmentation
}
// Set createdAt if it doesn't exist
if (!graphNoun.createdAt) {
const now = new Date()
graphNoun.createdAt = {
seconds: Math.floor(now.getTime() / 1000),
nanoseconds: (now.getTime() % 1000) * 1000000
}
}
}
// Always update the updatedAt timestamp
const now = new Date()
graphNoun.updatedAt = {
seconds: Math.floor(now.getTime() / 1000),
nanoseconds: (now.getTime() % 1000) * 1000000
}
}
// Update metadata
await this.storage!.saveMetadata(id, metadata)
// Track metadata statistics
const service = options.service || this.getCurrentAugmentation()
await this.storage!.incrementStatistic('metadata', service)
return true
} catch (error) {
console.error(`Failed to update metadata for vector ${id}:`, error)
throw new Error(`Failed to update metadata for vector ${id}: ${error}`)
}
}
/**
* 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<string> {
return this.addVerb(sourceId, targetId, undefined, {
type: relationType,
metadata: metadata
})
}
/**
* Create a connection between two entities
* This is an alias for relate() for backward compatibility
*/
public async connect(
sourceId: string,
targetId: string,
relationType: string,
metadata?: any
): Promise<string> {
return this.relate(sourceId, targetId, relationType, metadata)
}
/**
* Add a verb between two nouns
* If metadata is provided and vector is not, the metadata will be vectorized using the embedding function
*
* @param sourceId ID of the source noun
* @param targetId ID of the target noun
* @param vector Optional vector for the verb
* @param options Additional options:
* - type: Type of the verb
* - weight: Weight of the verb
* - metadata: Metadata for the verb
* - forceEmbed: Force using the embedding function for metadata even if vector is provided
* - id: Optional ID to use instead of generating a new one
* - autoCreateMissingNouns: Automatically create missing nouns if they don't exist
* - missingNounMetadata: Metadata to use when auto-creating missing nouns
*
* @returns The ID of the added verb
*
* @throws Error if source or target nouns don't exist and autoCreateMissingNouns is false or auto-creation fails
*/
public async addVerb(
sourceId: string,
targetId: string,
vector?: Vector,
options: {
type?: string
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
autoCreateMissingNouns?: boolean // Automatically create missing nouns
missingNounMetadata?: any // Metadata to use when auto-creating missing nouns
service?: string // The service that is inserting the data
} = {}
): Promise<string> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
// Check if source and target nouns exist
let sourceNoun = this.index.getNouns().get(sourceId)
let targetNoun = this.index.getNouns().get(targetId)
// Auto-create missing nouns if option is enabled
if (!sourceNoun && options.autoCreateMissingNouns) {
try {
// Create a placeholder vector for the missing noun
const placeholderVector = new Array(this._dimensions).fill(0)
// Add metadata if provided
const service = options.service || this.getCurrentAugmentation()
const now = new Date()
const timestamp = {
seconds: Math.floor(now.getTime() / 1000),
nanoseconds: (now.getTime() % 1000) * 1000000
}
const metadata = options.missingNounMetadata || {
autoCreated: true,
createdAt: timestamp,
updatedAt: timestamp,
noun: NounType.Concept,
createdBy: {
augmentation: service,
version: '1.0' // TODO: Get actual version from augmentation
}
}
// Add the missing noun
await this.add(placeholderVector, metadata, {id: sourceId})
// Get the newly created noun
sourceNoun = this.index.getNouns().get(sourceId)
console.warn(`Auto-created missing source noun with ID ${sourceId}`)
} catch (createError) {
console.error(`Failed to auto-create source noun with ID ${sourceId}:`, createError)
throw new Error(`Failed to auto-create source noun with ID ${sourceId}: ${createError}`)
}
}
if (!targetNoun && options.autoCreateMissingNouns) {
try {
// Create a placeholder vector for the missing noun
const placeholderVector = new Array(this._dimensions).fill(0)
// Add metadata if provided
const service = options.service || this.getCurrentAugmentation()
const now = new Date()
const timestamp = {
seconds: Math.floor(now.getTime() / 1000),
nanoseconds: (now.getTime() % 1000) * 1000000
}
const metadata = options.missingNounMetadata || {
autoCreated: true,
createdAt: timestamp,
updatedAt: timestamp,
noun: NounType.Concept,
createdBy: {
augmentation: service,
version: '1.0' // TODO: Get actual version from augmentation
}
}
// Add the missing noun
await this.add(placeholderVector, metadata, {id: targetId})
// Get the newly created noun
targetNoun = this.index.getNouns().get(targetId)
console.warn(`Auto-created missing target noun with ID ${targetId}`)
} catch (createError) {
console.error(`Failed to auto-create target noun with ID ${targetId}:`, createError)
throw new Error(`Failed to auto-create target noun with ID ${targetId}: ${createError}`)
}
}
if (!sourceNoun) {
throw new Error(`Source noun with ID ${sourceId} not found`)
}
if (!targetNoun) {
throw new Error(`Target noun with ID ${targetId} not found`)
}
// Use provided ID or generate a new one
const id = options.id || uuidv4()
let verbVector: Vector
// If metadata is provided and no vector is provided or forceEmbed is true, vectorize the metadata
if (options.metadata && (!vector || options.forceEmbed)) {
try {
// Extract a string representation from metadata for embedding
let textToEmbed: string
if (typeof options.metadata === 'string') {
textToEmbed = options.metadata
} else if (
options.metadata.description &&
typeof options.metadata.description === 'string'
) {
textToEmbed = options.metadata.description
} else {
// Convert to JSON string as fallback
textToEmbed = JSON.stringify(options.metadata)
}
// Ensure textToEmbed is a string
if (typeof textToEmbed !== 'string') {
textToEmbed = String(textToEmbed)
}
verbVector = await this.embeddingFunction(textToEmbed)
} catch (embedError) {
throw new Error(`Failed to vectorize verb metadata: ${embedError}`)
}
} else {
// Use a provided vector or average of source and target vectors
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 VerbType
)
if (!isValidVerbType) {
console.warn(
`Invalid verb type: ${verbType}. Using RelatedTo as default.`
)
// Set a default verb type
verbType = VerbType.RelatedTo
}
}
// Create verb
const verb: GraphVerb = {
id,
vector: verbVector,
connections: new Map(),
sourceId,
targetId,
type: verbType,
weight: options.weight,
metadata: options.metadata
}
// Add to index
await this.index.addItem({id, vector: verbVector})
// Get the noun from the index
const indexNoun = this.index.getNouns().get(id)
if (!indexNoun) {
throw new Error(
`Failed to retrieve newly created verb noun with ID ${id}`
)
}
// Update verb connections from index
verb.connections = indexNoun.connections
// Save verb to storage
await this.storage!.saveVerb(verb)
// Track verb statistics
const service = options.service || 'default'
await this.storage!.incrementStatistic('verb', service)
// Update HNSW index size (excluding verbs)
await this.storage!.updateHnswIndexSize(await this.getNounCount())
return id
} catch (error) {
console.error('Failed to add verb:', error)
throw new Error(`Failed to add verb: ${error}`)
}
}
/**
* Get a verb by ID
*/
public async getVerb(id: string): Promise<GraphVerb | null> {
await this.ensureInitialized()
try {
return await this.storage!.getVerb(id)
} catch (error) {
console.error(`Failed to get verb ${id}:`, error)
throw new Error(`Failed to get verb ${id}: ${error}`)
}
}
/**
* Get all verbs
*/
public async getAllVerbs(): Promise<GraphVerb[]> {
await this.ensureInitialized()
try {
return await this.storage!.getAllVerbs()
} catch (error) {
console.error('Failed to get all verbs:', error)
throw new Error(`Failed to get all verbs: ${error}`)
}
}
/**
* Get verbs by source noun ID
*/
public async getVerbsBySource(sourceId: string): Promise<GraphVerb[]> {
await this.ensureInitialized()
try {
return await this.storage!.getVerbsBySource(sourceId)
} catch (error) {
console.error(`Failed to get verbs by source ${sourceId}:`, error)
throw new Error(`Failed to get verbs by source ${sourceId}: ${error}`)
}
}
/**
* Get verbs by target noun ID
*/
public async getVerbsByTarget(targetId: string): Promise<GraphVerb[]> {
await this.ensureInitialized()
try {
return await this.storage!.getVerbsByTarget(targetId)
} catch (error) {
console.error(`Failed to get verbs by target ${targetId}:`, error)
throw new Error(`Failed to get verbs by target ${targetId}: ${error}`)
}
}
/**
* Get verbs by type
*/
public async getVerbsByType(type: string): Promise<GraphVerb[]> {
await this.ensureInitialized()
try {
return await this.storage!.getVerbsByType(type)
} catch (error) {
console.error(`Failed to get verbs by type ${type}:`, error)
throw new Error(`Failed to get verbs by type ${type}: ${error}`)
}
}
/**
* Delete a verb
* @param id The ID of the verb to delete
* @param options Additional options
* @returns Promise that resolves to true if the verb was deleted, false otherwise
*/
public async deleteVerb(
id: string,
options: {
service?: string // The service that is deleting the data
} = {}
): Promise<boolean> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
// Remove from index
const removed = this.index.removeItem(id)
if (!removed) {
return false
}
// Remove from storage
await this.storage!.deleteVerb(id)
// Track deletion statistics
const service = options.service || 'default'
await this.storage!.decrementStatistic('verb', service)
return true
} catch (error) {
console.error(`Failed to delete verb ${id}:`, error)
throw new Error(`Failed to delete verb ${id}: ${error}`)
}
}
/**
* Clear the database
*/
public async clear(): Promise<void> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
try {
// Clear index
this.index.clear()
// Clear storage
await this.storage!.clear()
} catch (error) {
console.error('Failed to clear vector database:', error)
throw new Error(`Failed to clear vector database: ${error}`)
}
}
/**
* Get the number of vectors in the database
*/
public size(): number {
return this.index.size()
}
/**
* Get the number of nouns in the database (excluding verbs)
* This is used for statistics reporting to match the expected behavior in tests
* @private
*/
private async getNounCount(): Promise<number> {
// Get all verbs from storage
const allVerbs = await this.storage!.getAllVerbs()
// Create a set of verb IDs for faster lookup
const verbIds = new Set(allVerbs.map(verb => verb.id))
// Get all nouns from the index
const nouns = this.index.getNouns()
// Count nouns that are not verbs
let nounCount = 0
for (const [id] of nouns.entries()) {
if (!verbIds.has(id)) {
nounCount++
}
}
return nounCount
}
/**
* Get statistics about the current state of the database
* @param options Additional options for retrieving statistics
* @returns Object containing counts of nouns, verbs, metadata entries, and HNSW index size
*/
public async getStatistics(options: {
service?: string | string[] // Filter statistics by service(s)
} = {}): Promise<{
nounCount: number
verbCount: number
metadataCount: number
hnswIndexSize: number
serviceBreakdown?: {
[service: string]: {
nounCount: number
verbCount: number
metadataCount: number
}
}
}> {
await this.ensureInitialized()
try {
// Get statistics from storage
const stats = await this.storage!.getStatistics()
// If statistics are available, use them
if (stats) {
// Initialize result
const result = {
nounCount: 0,
verbCount: 0,
metadataCount: 0,
hnswIndexSize: stats.hnswIndexSize,
serviceBreakdown: {} as {
[service: string]: {
nounCount: number
verbCount: number
metadataCount: number
}
}
}
// Filter by service if specified
const services = options.service
? (Array.isArray(options.service) ? options.service : [options.service])
: Object.keys({...stats.nounCount, ...stats.verbCount, ...stats.metadataCount})
// Calculate totals and service breakdown
for (const service of services) {
const nounCount = stats.nounCount[service] || 0
const verbCount = stats.verbCount[service] || 0
const metadataCount = stats.metadataCount[service] || 0
// Add to totals
result.nounCount += nounCount
result.verbCount += verbCount
result.metadataCount += metadataCount
// Add to service breakdown
result.serviceBreakdown[service] = {
nounCount,
verbCount,
metadataCount
}
}
return result
}
// If statistics are not available, fall back to calculating them on-demand
console.warn('Persistent statistics not available, calculating on-demand')
// Get all verbs from storage
const allVerbs = await this.storage!.getAllVerbs()
const verbCount = allVerbs.length
// Get the noun count using the helper method
const nounCount = await this.getNounCount()
// Count metadata entries by checking each noun for metadata
let metadataCount = 0
const nouns = this.index.getNouns()
for (const [id] of nouns.entries()) {
try {
const metadata = await this.storage!.getMetadata(id)
if (metadata !== null && metadata !== undefined) {
metadataCount++
}
} catch (error) {
// Ignore errors when checking individual metadata entries
// This could happen if metadata is corrupted or missing
}
}
// Get HNSW index size (excluding verbs)
// The HNSW index includes both nouns and verbs, but for statistics we want to report
// only the number of actual nouns (excluding verbs) to match the expected behavior in tests
const hnswIndexSize = nounCount
// Create default statistics
const defaultStats = {
nounCount,
verbCount,
metadataCount,
hnswIndexSize
}
// Initialize persistent statistics
const service = 'default'
await this.storage!.saveStatistics({
nounCount: { [service]: nounCount },
verbCount: { [service]: verbCount },
metadataCount: { [service]: metadataCount },
hnswIndexSize,
lastUpdated: new Date().toISOString()
})
return defaultStats
} catch (error) {
console.error('Failed to get statistics:', error)
throw new Error(`Failed to get statistics: ${error}`)
}
}
/**
* Check if the database is in read-only mode
* @returns True if the database is in read-only mode, false otherwise
*/
public isReadOnly(): boolean {
return this.readOnly
}
/**
* Set the database to read-only mode
* @param readOnly True to set the database to read-only mode, false to allow writes
*/
public setReadOnly(readOnly: boolean): void {
this.readOnly = readOnly
}
/**
* Embed text or data into a vector using the same embedding function used by this instance
* This allows clients to use the same TensorFlow Universal Sentence Encoder throughout their application
*
* @param data Text or data to embed
* @returns A promise that resolves to the embedded vector
*/
public async embed(data: string | string[]): Promise<Vector> {
await this.ensureInitialized()
try {
return await this.embeddingFunction(data)
} catch (error) {
console.error('Failed to embed data:', error)
throw new Error(`Failed to embed data: ${error}`)
}
}
/**
* Search for verbs by type and/or vector similarity
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of verbs with similarity scores
*/
public async searchVerbs(
queryVectorOrData: Vector | any,
k: number = 10,
options: {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
verbTypes?: string[] // Optional array of verb types to search within
service?: string // Filter results by the service that created the data
} = {}
): Promise<Array<GraphVerb & { similarity: number }>> {
await this.ensureInitialized()
try {
let queryVector: Vector
// Check if input is already a vector
if (
Array.isArray(queryVectorOrData) &&
queryVectorOrData.every((item) => typeof item === 'number') &&
!options.forceEmbed
) {
// Input is already a vector
queryVector = queryVectorOrData
} else {
// Input needs to be vectorized
try {
queryVector = await this.embeddingFunction(queryVectorOrData)
} catch (embedError) {
throw new Error(`Failed to vectorize query data: ${embedError}`)
}
}
// First use the HNSW index to find similar vectors efficiently
const searchResults = await this.index.search(queryVector, k * 2)
// Get all verbs for filtering
const allVerbs = await this.storage!.getAllVerbs()
// Create a map of verb IDs for faster lookup
const verbMap = new Map<string, GraphVerb>()
for (const verb of allVerbs) {
verbMap.set(verb.id, verb)
}
// Filter search results to only include verbs
const verbResults: Array<GraphVerb & { similarity: number }> = []
for (const result of searchResults) {
// Search results are [id, distance] tuples
const [id, distance] = result
const verb = verbMap.get(id)
if (verb) {
// If verb types are specified, check if this verb matches
if (options.verbTypes && options.verbTypes.length > 0) {
if (!verb.type || !options.verbTypes.includes(verb.type)) {
continue
}
}
verbResults.push({
...verb,
similarity: distance
})
}
}
// If we didn't get enough results from the index, fall back to the old method
if (verbResults.length < k) {
console.warn('Not enough verb results from HNSW index, falling back to manual search')
// Get verbs to search through
let verbs: GraphVerb[] = []
// If verb types are specified, get verbs of those types
if (options.verbTypes && options.verbTypes.length > 0) {
// Get verbs for each verb type in parallel
const verbPromises = options.verbTypes.map((verbType) =>
this.getVerbsByType(verbType)
)
const verbArrays = await Promise.all(verbPromises)
// Combine all verbs
for (const verbArray of verbArrays) {
verbs.push(...verbArray)
}
} else {
// Use all verbs
verbs = allVerbs
}
// Calculate similarity for each verb not already in results
const existingIds = new Set(verbResults.map(v => v.id))
for (const verb of verbs) {
if (!existingIds.has(verb.id) && verb.vector && verb.vector.length > 0) {
const distance = this.index.getDistanceFunction()(
queryVector,
verb.vector
)
verbResults.push({
...verb,
similarity: distance
})
}
}
}
// Sort by similarity (ascending distance)
verbResults.sort((a, b) => a.similarity - b.similarity)
// Take top k results
return verbResults.slice(0, k)
} catch (error) {
console.error('Failed to search verbs:', error)
throw new Error(`Failed to search verbs: ${error}`)
}
}
/**
* Search for nouns connected by specific verb types
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async searchNounsByVerbs(
queryVectorOrData: Vector | any,
k: number = 10,
options: {
forceEmbed?: boolean // Force using the embedding function even if input is a vector
verbTypes?: string[] // Optional array of verb types to filter by
direction?: 'outgoing' | 'incoming' | 'both' // Direction of verbs to consider
} = {}
): Promise<SearchResult<T>[]> {
await this.ensureInitialized()
try {
// First, search for nouns
const nounResults = await this.searchByNounTypes(
queryVectorOrData,
k * 2, // Get more results initially to account for filtering
null,
{forceEmbed: options.forceEmbed}
)
// If no verb types specified, return the noun results directly
if (!options.verbTypes || options.verbTypes.length === 0) {
return nounResults.slice(0, k)
}
// For each noun, get connected nouns through specified verb types
const connectedNounIds = new Set<string>()
const direction = options.direction || 'both'
for (const result of nounResults) {
// Get verbs connected to this noun
let connectedVerbs: GraphVerb[] = []
if (direction === 'outgoing' || direction === 'both') {
// Get outgoing verbs
const outgoingVerbs = await this.storage!.getVerbsBySource(result.id)
connectedVerbs.push(...outgoingVerbs)
}
if (direction === 'incoming' || direction === 'both') {
// Get incoming verbs
const incomingVerbs = await this.storage!.getVerbsByTarget(result.id)
connectedVerbs.push(...incomingVerbs)
}
// Filter by verb types if specified
if (options.verbTypes && options.verbTypes.length > 0) {
connectedVerbs = connectedVerbs.filter(
(verb) => verb.verb && options.verbTypes!.includes(verb.verb)
)
}
// Add connected noun IDs to the set
for (const verb of connectedVerbs) {
if (verb.source && verb.source !== result.id) {
connectedNounIds.add(verb.source)
}
if (verb.target && verb.target !== result.id) {
connectedNounIds.add(verb.target)
}
}
}
// Get the connected nouns
const connectedNouns: SearchResult<T>[] = []
for (const id of connectedNounIds) {
try {
const noun = this.index.getNouns().get(id)
if (noun) {
const metadata = await this.storage!.getMetadata(id)
// Calculate similarity score
let queryVector: Vector
if (
Array.isArray(queryVectorOrData) &&
queryVectorOrData.every((item) => typeof item === 'number') &&
!options.forceEmbed
) {
queryVector = queryVectorOrData
} else {
queryVector = await this.embeddingFunction(queryVectorOrData)
}
const distance = this.index.getDistanceFunction()(
queryVector,
noun.vector
)
connectedNouns.push({
id,
score: distance,
vector: noun.vector,
metadata: metadata as T | undefined
})
}
} catch (error) {
console.warn(`Failed to retrieve noun ${id}:`, error)
}
}
// Sort by similarity score
connectedNouns.sort((a, b) => a.score - b.score)
// Return top k results
return connectedNouns.slice(0, k)
} catch (error) {
console.error('Failed to search nouns by verbs:', error)
throw new Error(`Failed to search nouns by verbs: ${error}`)
}
}
/**
* Search for similar documents using a text query
* This is a convenience method that embeds the query text and performs a search
*
* @param query Text query to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async searchText(
query: string,
k: number = 10,
options: {
nounTypes?: string[]
includeVerbs?: boolean
searchMode?: 'local' | 'remote' | 'combined'
} = {}
): Promise<SearchResult<T>[]> {
await this.ensureInitialized()
try {
// Embed the query text
const queryVector = await this.embed(query)
// Search using the embedded vector
return await this.search(queryVector, k, {
nounTypes: options.nounTypes,
includeVerbs: options.includeVerbs,
searchMode: options.searchMode
})
} catch (error) {
console.error('Failed to search with text query:', error)
throw new Error(`Failed to search with text query: ${error}`)
}
}
/**
* Search a remote Brainy server for similar vectors
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async searchRemote(
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
storeResults?: boolean // Whether to store the results in the local database (default: true)
service?: string // Filter results by the service that created the data
} = {}
): Promise<SearchResult<T>[]> {
await this.ensureInitialized()
// Check if connected to a remote server
if (!this.isConnectedToRemoteServer()) {
throw new Error(
'Not connected to a remote server. Call connectToRemoteServer() first.'
)
}
try {
// If input is a string, convert it to a query string for the server
let query: string
if (typeof queryVectorOrData === 'string') {
query = queryVectorOrData
} else {
// For vectors, we need to embed them as a string query
// This is a simplification - ideally we would send the vector directly
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,
query,
k
)
if (!searchResult.success) {
throw new Error(`Remote search failed: ${searchResult.error}`)
}
return searchResult.data as SearchResult<T>[]
} catch (error) {
console.error('Failed to search remote server:', error)
throw new Error(`Failed to search remote server: ${error}`)
}
}
/**
* Search both local and remote Brainy instances, combining the results
* @param queryVectorOrData Query vector or data to search for
* @param k Number of results to return
* @param options Additional options
* @returns Array of search results
*/
public async searchCombined(
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
localFirst?: boolean // Whether to search local first (default: true)
service?: string // Filter results by the service that created the data
} = {}
): Promise<SearchResult<T>[]> {
await this.ensureInitialized()
// Check if connected to a remote server
if (!this.isConnectedToRemoteServer()) {
// If not connected to a remote server, just search locally
return this.searchLocal(queryVectorOrData, k, options)
}
try {
// Default to searching local first
const localFirst = options.localFirst !== false
if (localFirst) {
// Search local first
const localResults = await this.searchLocal(
queryVectorOrData,
k,
options
)
// If we have enough local results, return them
if (localResults.length >= k) {
return localResults
}
// Otherwise, search remote for additional results
const remoteResults = await this.searchRemote(
queryVectorOrData,
k - localResults.length,
{...options, storeResults: true}
)
// Combine results, removing duplicates
const combinedResults = [...localResults]
const localIds = new Set(localResults.map((r) => r.id))
for (const result of remoteResults) {
if (!localIds.has(result.id)) {
combinedResults.push(result)
}
}
return combinedResults
} else {
// Search remote first
const remoteResults = await this.searchRemote(queryVectorOrData, k, {
...options,
storeResults: true
})
// If we have enough remote results, return them
if (remoteResults.length >= k) {
return remoteResults
}
// Otherwise, search local for additional results
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))
for (const result of localResults) {
if (!remoteIds.has(result.id)) {
combinedResults.push(result)
}
}
return combinedResults
}
} catch (error) {
console.error('Failed to perform combined search:', error)
throw new Error(`Failed to perform combined search: ${error}`)
}
}
/**
* Check if the instance is connected to a remote server
* @returns True if connected to a remote server, false otherwise
*/
public isConnectedToRemoteServer(): boolean {
return !!(this.serverSearchConduit && this.serverConnection)
}
/**
* Disconnect from the remote server
* @returns True if successfully disconnected, false if not connected
*/
public async disconnectFromRemoteServer(): Promise<boolean> {
if (!this.isConnectedToRemoteServer()) {
return false
}
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
)
// Clear the connection information
this.serverSearchConduit = null
this.serverConnection = null
return true
} catch (error) {
console.error('Failed to disconnect from remote server:', error)
throw new Error(`Failed to disconnect from remote server: ${error}`)
}
}
/**
* Ensure the database is initialized
*/
private async ensureInitialized(): Promise<void> {
if (this.isInitialized) {
return
}
if (this.isInitializing) {
// If initialization is already in progress, wait for it to complete
// by polling the isInitialized flag
let attempts = 0
const maxAttempts = 100 // Prevent infinite loop
const delay = 50 // ms
while (
this.isInitializing &&
!this.isInitialized &&
attempts < maxAttempts
) {
await new Promise((resolve) => setTimeout(resolve, delay))
attempts++
}
if (!this.isInitialized) {
// If still not initialized after waiting, try to initialize again
await this.init()
}
} else {
// Normal case - not initialized and not initializing
await this.init()
}
}
/**
* Get information about the current storage usage and capacity
* @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<string, any>
}> {
await this.ensureInitialized()
if (!this.storage) {
return {
type: 'any',
used: 0,
quota: null,
details: {error: 'Storage not initialized'}
}
}
try {
// 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', '')
return {
type: storageType || 'any',
used: 0,
quota: null,
details: {
error: 'Storage adapter does not implement getStorageStatus method',
storageAdapter: this.storage.constructor.name,
indexSize: this.size()
}
}
}
// Get storage status from the storage adapter
const storageStatus = await this.storage.getStorageStatus()
// Add index information to the details
let indexInfo: Record<string, any> = {
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 || 'any',
used: storageStatus.used || 0,
quota: storageStatus.quota || null,
details: {
...(storageStatus.details || {}),
index: indexInfo
}
}
} catch (error) {
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', '')
return {
type: storageType || 'any',
used: 0,
quota: null,
details: {
error: String(error),
storageAdapter: this.storage.constructor.name,
indexSize: this.size()
}
}
}
}
/**
* Shut down the database and clean up resources
* This should be called when the database is no longer needed
*/
public async shutDown(): Promise<void> {
try {
// Disconnect from remote server if connected
if (this.isConnectedToRemoteServer()) {
await this.disconnectFromRemoteServer()
}
// Clean up worker pools to release resources
cleanupWorkerPools()
// Additional cleanup could be added here in the future
this.isInitialized = false
} catch (error) {
console.error('Failed to shut down BrainyData:', error)
throw new Error(`Failed to shut down BrainyData: ${error}`)
}
}
/**
* 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<T>[]
verbs: GraphVerb[]
nounTypes: string[]
verbTypes: string[]
version: string
hnswIndex?: {
entryPointId: string | null
maxLevel: number
dimension: number | null
config: HNSWConfig
connections: Record<string, Record<string, string[]>>
}
}> {
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<string, Record<string, string[]>>
}
// Convert Map<number, Set<string>> 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}`)
}
}
/**
* Import sparse data into the database
* @param data The sparse data to import
* If vectors are not present for nouns, they will be created using the embedding function
* @param options Import options
* @returns Object containing counts of imported items
*/
public async importSparseData(
data: {
nouns: VectorDocument<T>[]
verbs: GraphVerb[]
nounTypes?: string[]
verbTypes?: string[]
hnswIndex?: {
entryPointId: string | null
maxLevel: number
dimension: number | null
config: HNSWConfig
connections: Record<string, Record<string, string[]>>
}
version: string
},
options: {
clearExisting?: boolean
} = {}
): Promise<{
nounsRestored: number
verbsRestored: number
}> {
return this.restore(data, options)
}
/**
* 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<T>[]
verbs: GraphVerb[]
nounTypes?: string[]
verbTypes?: string[]
hnswIndex?: {
entryPointId: string | null
maxLevel: number
dimension: number | null
config: HNSWConfig
connections: Record<string, Record<string, string[]>>
}
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) {
await 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[]
}> {
await this.ensureInitialized()
// Check if database is in read-only mode
this.checkReadOnly()
// Set default options
const nounCount = options.nounCount || 10
const verbCount = options.verbCount || 20
const nounTypes = options.nounTypes || Object.values(NounType)
const verbTypes = options.verbTypes || Object.values(VerbType)
const clearExisting = options.clearExisting || false
// Clear existing data if requested
if (clearExisting) {
await this.clear()
}
try {
// Generate random nouns
const nounIds: string[] = []
const nounDescriptions: Record<string, string> = {
[NounType.Person]: 'A person with unique characteristics',
[NounType.Place]: 'A location with specific attributes',
[NounType.Thing]: 'An object with distinct properties',
[NounType.Event]: 'An occurrence with temporal aspects',
[NounType.Concept]: 'An abstract idea or notion',
[NounType.Content]: 'A piece of content or information',
[NounType.Group]: 'A collection of related entities',
[NounType.List]: 'An ordered sequence of items',
[NounType.Category]: 'A classification or grouping'
}
for (let i = 0; i < nounCount; i++) {
// Select a random noun type
const nounType = nounTypes[Math.floor(Math.random() * nounTypes.length)]
// Generate a random label
const label = `Random ${nounType} ${i + 1}`
// Create metadata
const metadata = {
noun: nounType,
label,
description: nounDescriptions[nounType] || `A random ${nounType}`,
randomAttributes: {
value: Math.random() * 100,
priority: Math.floor(Math.random() * 5) + 1,
tags: [`tag-${i % 5}`, `category-${i % 3}`]
}
}
// Add the noun
const id = await this.add(metadata.description, metadata as T)
nounIds.push(id)
}
// Generate random verbs between nouns
const verbIds: string[] = []
const verbDescriptions: Record<string, string> = {
[VerbType.AttributedTo]: 'Attribution relationship',
[VerbType.Controls]: 'Control relationship',
[VerbType.Created]: 'Creation relationship',
[VerbType.Earned]: 'Achievement relationship',
[VerbType.Owns]: 'Ownership relationship',
[VerbType.MemberOf]: 'Membership relationship',
[VerbType.RelatedTo]: 'General relationship',
[VerbType.WorksWith]: 'Collaboration relationship',
[VerbType.FriendOf]: 'Friendship relationship',
[VerbType.ReportsTo]: 'Reporting relationship',
[VerbType.Supervises]: 'Supervision relationship',
[VerbType.Mentors]: 'Mentorship relationship'
}
for (let i = 0; i < verbCount; i++) {
// Select random source and target nouns
const sourceIndex = Math.floor(Math.random() * nounIds.length)
let targetIndex = Math.floor(Math.random() * nounIds.length)
// Ensure source and target are different
while (targetIndex === sourceIndex && nounIds.length > 1) {
targetIndex = Math.floor(Math.random() * nounIds.length)
}
const sourceId = nounIds[sourceIndex]
const targetId = nounIds[targetIndex]
// Select a random verb type
const verbType = verbTypes[Math.floor(Math.random() * verbTypes.length)]
// Create metadata
const metadata = {
verb: verbType,
description:
verbDescriptions[verbType] || `A random ${verbType} relationship`,
weight: Math.random(),
confidence: Math.random(),
randomAttributes: {
strength: Math.random() * 100,
duration: Math.floor(Math.random() * 365) + 1,
tags: [`relation-${i % 5}`, `strength-${i % 3}`]
}
}
// Add the verb
const id = await this.addVerb(sourceId, targetId, undefined, {
type: verbType,
weight: metadata.weight,
metadata
})
verbIds.push(id)
}
return {
nounIds,
verbIds
}
} catch (error) {
console.error('Failed to generate random graph:', error)
throw new Error(`Failed to generate random graph: ${error}`)
}
}
}
// Export distance functions for convenience
export {
euclideanDistance,
cosineDistance,
manhattanDistance,
dotProductDistance
} from './utils/index.js'
Reference in a new issue View git blame Copy permalink