import { AugmentationType, IMemoryAugmentation, AugmentationResponse } from '../types/augmentations.js' import {StorageAdapter, Vector} from '../coreTypes.js' import {MemoryStorage, OPFSStorage} from '../storage/storageFactory.js' // FileSystemStorage will be dynamically imported when needed to avoid fs imports in browser import {cosineDistance} from '../utils/distance.js' /** * Base class for memory augmentations that wrap a StorageAdapter */ abstract class BaseMemoryAugmentation implements IMemoryAugmentation { readonly name: string readonly description: string = 'Base memory augmentation' enabled: boolean = true protected storage: StorageAdapter protected isInitialized = false constructor(name: string, storage: StorageAdapter) { this.name = name this.storage = storage } async initialize(): Promise { if (this.isInitialized) { return } try { await this.storage.init() this.isInitialized = true } catch (error) { console.error(`Failed to initialize ${this.name}:`, error) throw new Error(`Failed to initialize ${this.name}: ${error}`) } } async shutDown(): Promise { this.isInitialized = false } async getStatus(): Promise<'active' | 'inactive' | 'error'> { return this.isInitialized ? 'active' : 'inactive' } async storeData( key: string, data: unknown, options?: Record ): Promise> { await this.ensureInitialized() try { await this.storage.saveMetadata(key, data) return {success: true, data: true} } catch (error) { console.error(`Failed to store data for key ${key}:`, error) return { success: false, data: false, error: `Failed to store data: ${error}` } } } async retrieveData( key: string, options?: Record ): Promise> { await this.ensureInitialized() try { const data = await this.storage.getMetadata(key) return { success: true, data } } catch (error) { console.error(`Failed to retrieve data for key ${key}:`, error) return { success: false, data: null, error: `Failed to retrieve data: ${error}` } } } async updateData( key: string, data: unknown, options?: Record ): Promise> { await this.ensureInitialized() try { await this.storage.saveMetadata(key, data) return {success: true, data: true} } catch (error) { console.error(`Failed to update data for key ${key}:`, error) return { success: false, data: false, error: `Failed to update data: ${error}` } } } async deleteData( key: string, options?: Record ): Promise> { await this.ensureInitialized() try { // There's no direct deleteMetadata method, so we save null await this.storage.saveMetadata(key, null) return {success: true, data: true} } catch (error) { console.error(`Failed to delete data for key ${key}:`, error) return { success: false, data: false, error: `Failed to delete data: ${error}` } } } async listDataKeys( pattern?: string, options?: Record ): Promise> { // This is a limitation of the current StorageAdapter interface // It doesn't provide a way to list all metadata keys // We could implement this in the future by extending the StorageAdapter interface return { success: false, data: [], error: 'listDataKeys is not supported by this storage adapter' } } /** * Searches for data in the storage using vector similarity. * Implements the findNearest functionality by calculating distances client-side. * @param query The query vector or data to search for * @param k Number of results to return (default: 10) * @param options Optional search options */ async search( query: unknown, k: number = 10, options?: Record ): Promise>> { await this.ensureInitialized() try { // Check if query is a vector let queryVector: Vector if (Array.isArray(query) && query.every(item => typeof item === 'number')) { queryVector = query as Vector } else { // If query is not a vector, we can't perform vector search return { success: false, data: [], error: 'Query must be a vector (array of numbers) for vector search' } } // Get all nodes from storage const nodes = await this.storage.getAllNouns() // Calculate distances and prepare results const results: Array<{ id: string; score: number; data: unknown; }> = [] for (const node of nodes) { // Skip nodes that don't have a vector if (!node.vector || !Array.isArray(node.vector)) { continue } // Get metadata for the node const metadata = await this.storage.getMetadata(node.id) // Calculate distance between query vector and node vector const distance = cosineDistance(queryVector, node.vector) // Convert distance to similarity score (1 - distance for cosine) // This way higher scores are better (more similar) const score = 1 - distance results.push({ id: node.id, score, data: metadata }) } // Sort results by score (descending) and take top k results.sort((a, b) => b.score - a.score) const topResults = results.slice(0, k) return { success: true, data: topResults } } catch (error) { console.error(`Failed to search in storage:`, error) return { success: false, data: [], error: `Failed to search in storage: ${error}` } } } protected async ensureInitialized(): Promise { if (!this.isInitialized) { await this.initialize() } } } /** * Memory augmentation that uses in-memory storage */ export class MemoryStorageAugmentation extends BaseMemoryAugmentation { readonly description = 'Memory augmentation that stores data in memory' enabled = true constructor(name: string) { super(name, new MemoryStorage()) } getType(): AugmentationType { return AugmentationType.MEMORY } } /** * Memory augmentation that uses file system storage */ export class FileSystemStorageAugmentation extends BaseMemoryAugmentation { readonly description = 'Memory augmentation that stores data in the file system' enabled = true private rootDirectory: string constructor(name: string, rootDirectory?: string) { // Temporarily use MemoryStorage, will be replaced in initialize() super(name, new MemoryStorage()) this.rootDirectory = rootDirectory || '.' } async initialize(): Promise { try { // Dynamically import FileSystemStorage const { FileSystemStorage } = await import('../storage/adapters/fileSystemStorage.js') this.storage = new FileSystemStorage(this.rootDirectory) await super.initialize() } catch (error) { console.error('Failed to load FileSystemStorage:', error) throw new Error(`Failed to initialize FileSystemStorage: ${error}`) } } getType(): AugmentationType { return AugmentationType.MEMORY } } /** * Memory augmentation that uses OPFS (Origin Private File System) storage */ export class OPFSStorageAugmentation extends BaseMemoryAugmentation { readonly description = 'Memory augmentation that stores data in the Origin Private File System' enabled = true constructor(name: string) { super(name, new OPFSStorage()) } getType(): AugmentationType { return AugmentationType.MEMORY } } /** * Factory function to create the appropriate memory augmentation based on the environment */ export async function createMemoryAugmentation( name: string, options: { storageType?: 'memory' | 'filesystem' | 'opfs' rootDirectory?: string requestPersistentStorage?: boolean } = {} ): Promise { // If a specific storage type is requested, use that if (options.storageType) { switch (options.storageType) { case 'memory': return new MemoryStorageAugmentation(name) case 'filesystem': return new FileSystemStorageAugmentation(name, options.rootDirectory) case 'opfs': return new OPFSStorageAugmentation(name) } } // Otherwise, select based on environment // Use the global isNode variable from the environment detection const isNodeEnv = (globalThis as any).__ENV__?.isNode || ( typeof process !== 'undefined' && process.versions != null && process.versions.node != null ) if (isNodeEnv) { // In Node.js, use FileSystemStorage return new FileSystemStorageAugmentation(name, options.rootDirectory) } else { // In browser, try OPFS first const opfsStorage = new OPFSStorage() if (opfsStorage.isOPFSAvailable()) { // Request persistent storage if specified if (options.requestPersistentStorage) { await opfsStorage.requestPersistentStorage() } return new OPFSStorageAugmentation(name) } else { // Fall back to memory storage return new MemoryStorageAugmentation(name) } } }