### 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.
1938 lines
57 KiB
TypeScript
1938 lines
57 KiB
TypeScript
/**
|
|
* 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/opfsStorage.js'
|
|
import {
|
|
DistanceFunction,
|
|
GraphVerb,
|
|
EmbeddingFunction,
|
|
HNSWConfig,
|
|
HNSWNoun,
|
|
SearchResult,
|
|
StorageAdapter,
|
|
Vector,
|
|
VectorDocument
|
|
} from './coreTypes.js'
|
|
import {
|
|
cosineDistance,
|
|
defaultEmbeddingFunction,
|
|
euclideanDistance
|
|
} from './utils/index.js'
|
|
import { NounType, VerbType, GraphNoun } from './types/graphTypes.js'
|
|
import {
|
|
ServerSearchConduitAugmentation,
|
|
createServerSearchAugmentations
|
|
} from './augmentations/serverSearchAugmentations.js'
|
|
import { WebSocketConnection } from './types/augmentations.js'
|
|
import { BrainyDataInterface } from './types/brainyDataInterface.js'
|
|
|
|
export interface BrainyDataConfig {
|
|
/**
|
|
* 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
|
|
|
|
/**
|
|
* Request persistent storage when running in a browser
|
|
* This will prompt the user for permission to use persistent storage
|
|
* @deprecated Use storage.requestPersistentStorage instead
|
|
*/
|
|
requestPersistentStorage?: boolean
|
|
|
|
/**
|
|
* 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
|
|
}
|
|
}
|
|
|
|
export class BrainyData<T = any> implements BrainyDataInterface<T> {
|
|
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: ServerSearchConduitAugmentation | null = null
|
|
private serverConnection: WebSocketConnection | null = null
|
|
|
|
/**
|
|
* Create a new vector database
|
|
*/
|
|
constructor(config: BrainyDataConfig = {}) {
|
|
// 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
|
|
|
|
// Set persistent storage request flag (support both new and deprecated options)
|
|
this.requestPersistentStorage =
|
|
config.storage?.requestPersistentStorage !== undefined
|
|
? config.storage.requestPersistentStorage
|
|
: config.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'
|
|
)
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Initialize the database
|
|
* Loads existing data from storage if available
|
|
*/
|
|
public async init(): Promise<void> {
|
|
if (this.isInitialized) {
|
|
return
|
|
}
|
|
|
|
try {
|
|
// Initialize storage if not provided in constructor
|
|
if (!this.storage) {
|
|
// Combine storage config with requestPersistentStorage for backward compatibility
|
|
const storageOptions = {
|
|
...this.storageConfig,
|
|
requestPersistentStorage: this.requestPersistentStorage
|
|
}
|
|
|
|
this.storage = await createStorage(storageOptions)
|
|
}
|
|
|
|
// 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) {
|
|
// Add to index
|
|
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
|
|
} catch (error) {
|
|
console.error('Failed to initialize BrainyData:', error)
|
|
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
|
|
} = {}
|
|
): 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')
|
|
}
|
|
|
|
// 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.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)
|
|
|
|
// 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
|
|
}
|
|
}
|
|
|
|
await this.storage!.saveMetadata(id, metadata)
|
|
}
|
|
|
|
// 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 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
|
|
} = {}
|
|
): Promise<string[]> {
|
|
await this.ensureInitialized()
|
|
|
|
// Check if database is in read-only mode
|
|
this.checkReadOnly()
|
|
|
|
const ids: string[] = []
|
|
|
|
try {
|
|
for (const item of items) {
|
|
const id = await this.add(item.vectorOrData, item.metadata, options)
|
|
ids.push(id)
|
|
}
|
|
|
|
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
|
|
} = {}
|
|
): 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 })
|
|
}
|
|
|
|
/**
|
|
* 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
|
|
} = {}
|
|
): Promise<SearchResult<T>[]> {
|
|
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}`)
|
|
}
|
|
}
|
|
|
|
// 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 = 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
|
|
}
|
|
|
|
const metadata = await this.storage!.getMetadata(id)
|
|
|
|
searchResults.push({
|
|
id,
|
|
score,
|
|
vector: noun.vector,
|
|
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 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
|
|
}
|
|
|
|
const metadata = await this.storage!.getMetadata(id)
|
|
|
|
searchResults.push({
|
|
id,
|
|
score,
|
|
vector: noun.vector,
|
|
metadata: metadata as T | undefined
|
|
})
|
|
}
|
|
|
|
return searchResults
|
|
}
|
|
} 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
|
|
} = {}
|
|
): Promise<SearchResult<T>[]> {
|
|
// 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
|
|
} = {}
|
|
): Promise<SearchResult<T>[]> {
|
|
// 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
|
|
}
|
|
)
|
|
} else {
|
|
// Otherwise, search all GraphNouns
|
|
searchResults = await this.searchByNounTypes(queryToUse, k, null, {
|
|
forceEmbed: options.forceEmbed
|
|
})
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
/**
|
|
* 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
|
|
*/
|
|
public async delete(id: string): 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)
|
|
|
|
// Try to remove metadata (ignore errors)
|
|
try {
|
|
await this.storage!.saveMetadata(id, null)
|
|
} 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
|
|
*/
|
|
public async updateMetadata(id: string, metadata: T): 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
|
|
}
|
|
}
|
|
|
|
// Update metadata
|
|
await this.storage!.saveMetadata(id, metadata)
|
|
|
|
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
|
|
})
|
|
}
|
|
|
|
/**
|
|
* Add a verb between two nouns
|
|
* If metadata is provided and vector is not, the metadata will be vectorized using the embedding function
|
|
*/
|
|
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
|
|
} = {}
|
|
): Promise<string> {
|
|
await this.ensureInitialized()
|
|
|
|
// Check if database is in read-only mode
|
|
this.checkReadOnly()
|
|
|
|
try {
|
|
// Check if source and target nouns exist
|
|
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`)
|
|
}
|
|
|
|
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
|
|
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)
|
|
|
|
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
|
|
*/
|
|
public async deleteVerb(id: string): 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)
|
|
|
|
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()
|
|
}
|
|
|
|
/**
|
|
* 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 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)
|
|
} = {}
|
|
): 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)
|
|
} = {}
|
|
): 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) {
|
|
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()
|
|
}
|
|
|
|
// 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}`)
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 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) {
|
|
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'
|