Initial commit of Brainy vector database v0.1.0

src/brainyData.ts

@@ -0,0 +1,549 @@
+/**
+ * BrainyData
+ * Main class that provides the vector database functionality
+ */
+
+import { v4 as uuidv4 } from 'uuid'
+import { HNSWIndex } from './hnsw/hnswIndex.js'
+import { createStorage } from './storage/opfsStorage.js'
+import { DistanceFunction, Edge, EmbeddingFunction, HNSWConfig, SearchResult, StorageAdapter, Vector, VectorDocument } from './coreTypes.js'
+import { cosineDistance, defaultEmbeddingFunction, euclideanDistance } from './utils/index.js'
+
+export interface BrainyDataConfig {
+  /**
+   * HNSW index configuration
+   */
+  hnsw?: Partial<HNSWConfig>
+
+  /**
+   * Distance function to use for similarity calculations
+   */
+  distanceFunction?: DistanceFunction
+
+  /**
+   * Custom storage adapter (if not provided, will use OPFS or memory storage)
+   */
+  storageAdapter?: StorageAdapter
+
+  /**
+   * Embedding function to convert data to vectors
+   */
+  embeddingFunction?: EmbeddingFunction
+}
+
+export class BrainyData<T = any> {
+  private index: HNSWIndex
+  private storage: StorageAdapter | null = null
+  private isInitialized = false
+  private embeddingFunction: EmbeddingFunction
+
+  /**
+   * Create a new vector database
+   */
+  constructor(config: BrainyDataConfig = {}) {
+    // Initialize HNSW index
+    this.index = new HNSWIndex(
+      config.hnsw,
+      config.distanceFunction || cosineDistance
+    )
+
+    // 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
+  }
+
+  /**
+   * 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) {
+        this.storage = await createStorage()
+      }
+
+      // Initialize storage
+      await this.storage!.init()
+
+      // Load all nodes from storage
+      const nodes = await this.storage!.getAllNodes()
+
+      // Clear the index and add all nodes
+      this.index.clear()
+      for (const node of nodes) {
+        // Add to index
+        this.index.addItem({
+          id: node.id,
+          vector: node.vector
+        })
+      }
+
+      this.isInitialized = true
+    } catch (error) {
+      console.error('Failed to initialize vector database:', error)
+      throw new Error(`Failed to initialize vector database: ${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
+    } = {}
+  ): Promise<string> {
+    await this.ensureInitialized()
+
+    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')
+      }
+
+      // Generate ID if isn't provided
+      const id = uuidv4()
+
+      // Add to index
+      this.index.addItem({ id, vector })
+
+      // Get the node from the index
+      const node = this.index.getNodes().get(id)
+
+      if (!node) {
+        throw new Error(`Failed to retrieve newly created node with ID ${id}`)
+      }
+
+      // Save node to storage
+      await this.storage!.saveNode(node)
+
+      // Save metadata if provided
+      if (metadata !== undefined) {
+        await this.storage!.saveMetadata(id, metadata)
+      }
+
+      return id
+    } catch (error) {
+      console.error('Failed to add vector:', error)
+      throw new Error(`Failed to add vector: ${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
+    } = {}
+  ): Promise<string[]> {
+    await this.ensureInitialized()
+
+    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}`)
+    }
+  }
+
+  /**
+   * 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
+    } = {}
+  ): 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')
+      }
+
+      // 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 node = this.index.getNodes().get(id)
+        if (!node) {
+          continue
+        }
+
+        const metadata = await this.storage!.getMetadata(id)
+
+        searchResults.push({
+          id,
+          score,
+          vector: node.vector,
+          metadata
+        })
+      }
+
+      return searchResults
+    } catch (error) {
+      console.error('Failed to search vectors:', error)
+      throw new Error(`Failed to search vectors: ${error}`)
+    }
+  }
+
+  /**
+   * Get a vector by ID
+   */
+  public async get(id: string): Promise<VectorDocument<T> | null> {
+    await this.ensureInitialized()
+
+    try {
+      // Get node from index
+      const node = this.index.getNodes().get(id)
+      if (!node) {
+        return null
+      }
+
+      // Get metadata
+      const metadata = await this.storage!.getMetadata(id)
+
+      return {
+        id,
+        vector: node.vector,
+        metadata
+      }
+    } catch (error) {
+      console.error(`Failed to get vector ${id}:`, error)
+      throw new Error(`Failed to get vector ${id}: ${error}`)
+    }
+  }
+
+  /**
+   * Delete a vector by ID
+   */
+  public async delete(id: string): Promise<boolean> {
+    await this.ensureInitialized()
+
+    try {
+      // Remove from index
+      const removed = this.index.removeItem(id)
+      if (!removed) {
+        return false
+      }
+
+      // Remove from storage
+      await this.storage!.deleteNode(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()
+
+    try {
+      // Check if a vector exists
+      const node = this.index.getNodes().get(id)
+      if (!node) {
+        return false
+      }
+
+      // 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}`)
+    }
+  }
+
+  /**
+   * Add an edge between two nodes
+   */
+  public async addEdge(
+    sourceId: string,
+    targetId: string,
+    vector?: Vector,
+    options: {
+      type?: string
+      weight?: number
+      metadata?: any
+    } = {}
+  ): Promise<string> {
+    await this.ensureInitialized()
+
+    try {
+      // Check if source and target nodes exist
+      const sourceNode = this.index.getNodes().get(sourceId)
+      const targetNode = this.index.getNodes().get(targetId)
+
+      if (!sourceNode) {
+        throw new Error(`Source node with ID ${sourceId} not found`)
+      }
+
+      if (!targetNode) {
+        throw new Error(`Target node with ID ${targetId} not found`)
+      }
+
+      // Generate ID for the edge
+      const id = uuidv4()
+
+      // Use a provided vector or average of source and target vectors
+      const edgeVector =
+        vector ||
+        sourceNode.vector.map((val, i) => (val + targetNode.vector[i]) / 2)
+
+      // Create edge
+      const edge: Edge = {
+        id,
+        vector: edgeVector,
+        connections: new Map(),
+        sourceId,
+        targetId,
+        type: options.type,
+        weight: options.weight,
+        metadata: options.metadata
+      }
+
+      // Add to index
+      this.index.addItem({ id, vector: edgeVector })
+
+      // Get the node from the index
+      const indexNode = this.index.getNodes().get(id)
+
+      if (!indexNode) {
+        throw new Error(
+          `Failed to retrieve newly created edge node with ID ${id}`
+        )
+      }
+
+      // Update edge connections from index
+      edge.connections = indexNode.connections
+
+      // Save edge to storage
+      await this.storage!.saveEdge(edge)
+
+      return id
+    } catch (error) {
+      console.error('Failed to add edge:', error)
+      throw new Error(`Failed to add edge: ${error}`)
+    }
+  }
+
+  /**
+   * Get an edge by ID
+   */
+  public async getEdge(id: string): Promise<Edge | null> {
+    await this.ensureInitialized()
+
+    try {
+      return await this.storage!.getEdge(id)
+    } catch (error) {
+      console.error(`Failed to get edge ${id}:`, error)
+      throw new Error(`Failed to get edge ${id}: ${error}`)
+    }
+  }
+
+  /**
+   * Get all edges
+   */
+  public async getAllEdges(): Promise<Edge[]> {
+    await this.ensureInitialized()
+
+    try {
+      return await this.storage!.getAllEdges()
+    } catch (error) {
+      console.error('Failed to get all edges:', error)
+      throw new Error(`Failed to get all edges: ${error}`)
+    }
+  }
+
+  /**
+   * Get edges by source node ID
+   */
+  public async getEdgesBySource(sourceId: string): Promise<Edge[]> {
+    await this.ensureInitialized()
+
+    try {
+      return await this.storage!.getEdgesBySource(sourceId)
+    } catch (error) {
+      console.error(`Failed to get edges by source ${sourceId}:`, error)
+      throw new Error(`Failed to get edges by source ${sourceId}: ${error}`)
+    }
+  }
+
+  /**
+   * Get edges by target node ID
+   */
+  public async getEdgesByTarget(targetId: string): Promise<Edge[]> {
+    await this.ensureInitialized()
+
+    try {
+      return await this.storage!.getEdgesByTarget(targetId)
+    } catch (error) {
+      console.error(`Failed to get edges by target ${targetId}:`, error)
+      throw new Error(`Failed to get edges by target ${targetId}: ${error}`)
+    }
+  }
+
+  /**
+   * Get edges by type
+   */
+  public async getEdgesByType(type: string): Promise<Edge[]> {
+    await this.ensureInitialized()
+
+    try {
+      return await this.storage!.getEdgesByType(type)
+    } catch (error) {
+      console.error(`Failed to get edges by type ${type}:`, error)
+      throw new Error(`Failed to get edges by type ${type}: ${error}`)
+    }
+  }
+
+  /**
+   * Delete an edge
+   */
+  public async deleteEdge(id: string): Promise<boolean> {
+    await this.ensureInitialized()
+
+    try {
+      // Remove from index
+      const removed = this.index.removeItem(id)
+      if (!removed) {
+        return false
+      }
+
+      // Remove from storage
+      await this.storage!.deleteEdge(id)
+
+      return true
+    } catch (error) {
+      console.error(`Failed to delete edge ${id}:`, error)
+      throw new Error(`Failed to delete edge ${id}: ${error}`)
+    }
+  }
+
+  /**
+   * Clear the database
+   */
+  public async clear(): Promise<void> {
+    await this.ensureInitialized()
+
+    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()
+  }
+
+  /**
+   * Ensure the database is initialized
+   */
+  private async ensureInitialized(): Promise<void> {
+    if (!this.isInitialized) {
+      await this.init()
+    }
+  }
+}
+
+// Export distance functions for convenience
+export { euclideanDistance, cosineDistance, manhattanDistance, dotProductDistance } from './utils/index.js'