personal-data-wallet-sdk/src/retrieval/MemoryRetrievalService.ts

TypeScript SDK for Personal Data Wallet - Decentralized memory system with AI embeddings, HNSW vector search, SEAL encryption and Walrus storage

/**
 * MemoryRetrievalService - Unified Memory Retrieval & Search
 * 
 * Consolidates all memory retrieval operations into a single, powerful service
 * with advanced caching, decryption, analytics, and multi-dimensional search.
 * 
 * Features:
 * - 🔍 Vector similarity search with HNSW indexing
 * - 📊 Semantic search with AI understanding
 * - 🕒 Time-based and metadata filtering
 * - 🔐 Automatic SEAL decryption pipeline
 * - 📈 Memory analytics and relationship discovery
 * - ⚡ Multi-tier caching for performance
 * - 🔄 Real-time memory streaming
 * - 📤 Export and backup capabilities
 */

import { EmbeddingService } from '../services/EmbeddingService';
import { VectorManager } from '../vector/VectorManager';
import { KnowledgeGraphManager } from '../graph/KnowledgeGraphManager';
import { StorageManager } from '../infrastructure/walrus/StorageManager';
import { BlockchainManager } from '../infrastructure/sui/BlockchainManager';
import { EncryptionService } from '../services/EncryptionService';
import { BatchManager } from '../batch/BatchManager';
import { MemoryDecryptionPipeline, DecryptionConfig } from './MemoryDecryptionPipeline';

// Enhanced types for unified retrieval
export interface UnifiedMemoryQuery {
  // Core search parameters
  query?: string;
  userId: string;
  
  // Search types
  searchType?: 'vector' | 'semantic' | 'keyword' | 'hybrid' | 'graph' | 'temporal';
  
  // Filtering options
  categories?: string[];
  dateRange?: {
    start?: Date;
    end?: Date;
  };
  importanceRange?: {
    min?: number;
    max?: number;
  };
  tags?: string[];
  contentTypes?: string[];
  
  // Vector search parameters
  similarity?: {
    threshold?: number;
    k?: number;
    efSearch?: number;
  };
  
  // Graph search parameters
  graphTraversal?: {
    maxDepth?: number;
    includeEntities?: string[];
    excludeRelations?: string[];
  };
  
  // Result options
  includeContent?: boolean;
  includeMetadata?: boolean;
  includeEmbeddings?: boolean;
  includeAnalytics?: boolean;
  
  // Performance options
  useCache?: boolean;
  timeout?: number;
  batchSize?: number;
}

export interface UnifiedMemoryResult {
  id: string;
  content?: string;
  category: string;
  created: Date;
  updated?: Date;
  
  // Scoring and relevance
  similarity?: number;
  relevanceScore: number;
  searchRelevance: {
    vectorSimilarity?: number;
    semanticMatch?: number;
    keywordMatch?: number;
    graphRelevance?: number;
    temporalRelevance?: number;
  };
  
  // Rich metadata
  metadata: {
    owner: string;
    size: number;
    importance: number;
    tags: string[];
    contentType: string;
    isEncrypted: boolean;
    accessCount?: number;
    lastAccessed?: Date;
  };
  
  // Storage information
  storage: {
    blobId: string;
    walrusHash?: string;
    blockchainId?: string;
    cacheStatus: 'cached' | 'retrieved' | 'not-cached';
  };
  
  // Relationships and context
  relationships?: {
    relatedMemories: string[];
    knowledgeGraphNodes: string[];
    semanticClusters: string[];
  };
  
  // Analytics
  analytics?: {
    viewCount: number;
    shareCount: number;
    editCount: number;
    sentimentScore?: number;
    topicDistribution?: Record<string, number>;
  };
}

export interface RetrievalStats {
  totalResults: number;
  processingTime: number;
  cacheHitRate: number;
  searchBreakdown: {
    vectorSearchTime: number;
    semanticAnalysisTime: number;
    decryptionTime: number;
    graphTraversalTime: number;
  };
  dataSourcesUsed: string[];
  qualityMetrics: {
    averageRelevance: number;
    diversityScore: number;
    freshnessScore: number;
  };
}

export interface RetrievalContext {
  query: UnifiedMemoryQuery;
  results: UnifiedMemoryResult[];
  stats: RetrievalStats;
  suggestions: {
    relatedQueries: string[];
    filterSuggestions: string[];
    explorationPaths: string[];
  };
  timeline?: {
    events: Array<{
      date: Date;
      type: 'created' | 'modified' | 'accessed' | 'shared';
      memoryId: string;
      description: string;
    }>;
  };
}

/**
 * Unified Memory Retrieval Service
 */
export class MemoryRetrievalService {
  private embeddingService: EmbeddingService;
  private vectorManager: VectorManager;
  private graphManager: KnowledgeGraphManager;
  private storageManager: StorageManager;
  private blockchainManager: BlockchainManager;
  private encryptionService: EncryptionService;
  private batchManager: BatchManager;
  private decryptionPipeline: MemoryDecryptionPipeline;
  
  // Multi-tier caching system
  private queryCache = new Map<string, { result: RetrievalContext; timestamp: number }>();
  private contentCache = new Map<string, { content: string; metadata: any; timestamp: number }>();
  private analyticsCache = new Map<string, { analytics: any; timestamp: number }>();
  
  // Cache TTL settings
  private readonly QUERY_CACHE_TTL = 5 * 60 * 1000; // 5 minutes
  private readonly CONTENT_CACHE_TTL = 30 * 60 * 1000; // 30 minutes
  private readonly ANALYTICS_CACHE_TTL = 60 * 60 * 1000; // 1 hour

  constructor(config?: {
    embeddingService?: EmbeddingService;
    vectorManager?: VectorManager;
    graphManager?: KnowledgeGraphManager;
    storageManager?: StorageManager;
    blockchainManager?: BlockchainManager;
    encryptionService?: EncryptionService;
    batchManager?: BatchManager;
    decryptionConfig?: DecryptionConfig;
  }) {
    // Initialize services (can be injected or created with default configs)
    this.embeddingService = config?.embeddingService ?? new EmbeddingService();
    this.storageManager = config?.storageManager ?? new StorageManager();
    this.vectorManager = config?.vectorManager ?? new VectorManager(
      this.storageManager as any, // TODO: Fix type compatibility
      {
        embedding: { apiKey: '' },
        index: { dimension: 768 },
        batch: { maxBatchSize: 10 }
      }
    );
    this.graphManager = config?.graphManager ?? new KnowledgeGraphManager();
    this.blockchainManager = config?.blockchainManager ?? new BlockchainManager();
    this.encryptionService = config?.encryptionService ?? new EncryptionService({} as any, {} as any);
    this.batchManager = config?.batchManager ?? new BatchManager();
    
    // Initialize decryption pipeline
    this.decryptionPipeline = new MemoryDecryptionPipeline(
      this.encryptionService,
      this.storageManager,
      config?.decryptionConfig
    );
  }

  // ==================== UNIFIED SEARCH ====================

  /**
   * Main unified search method - handles all types of memory retrieval
   */
  async searchMemories(query: UnifiedMemoryQuery): Promise<RetrievalContext> {
    const startTime = Date.now();
    
    // Check cache first
    const cacheKey = this.generateCacheKey(query);
    if (query.useCache !== false) {
      const cached = this.getCachedQuery(cacheKey);
      if (cached) {
        return cached;
      }
    }

    try {
      // Execute search based on type
      let results: UnifiedMemoryResult[];
      let searchStats: Partial<RetrievalStats['searchBreakdown']> = {};

      switch (query.searchType || 'hybrid') {
        case 'vector':
          results = await this.performVectorSearch(query, searchStats);
          break;
        case 'semantic':
          results = await this.performSemanticSearch(query, searchStats);
          break;
        case 'keyword':
          results = await this.performKeywordSearch(query, searchStats);
          break;
        case 'graph':
          results = await this.performGraphSearch(query, searchStats);
          break;
        case 'temporal':
          results = await this.performTemporalSearch(query, searchStats);
          break;
        case 'hybrid':
        default:
          results = await this.performHybridSearch(query, searchStats);
          break;
      }

      // Apply filters and post-processing
      results = await this.applyFilters(results, query);
      results = await this.enrichResults(results, query);
      results = this.rankAndSortResults(results, query);
      
      // Auto-decrypt encrypted memories if decryption is available
      if (this.decryptionPipeline.isReady()) {
        try {
          results = await this.decryptionPipeline.decryptMemoryResults(results, query.userId);
        } catch (error) {
          console.warn('⚠️ Auto-decryption failed, returning encrypted results:', error);
        }
      }

      // Generate stats and suggestions
      const stats = this.generateStats(results, searchStats, startTime);
      const suggestions = await this.generateSuggestions(query, results);
      const timeline = await this.generateTimeline(results, query);

      const context: RetrievalContext = {
        query,
        results,
        stats,
        suggestions,
        timeline
      };

      // Cache the results
      if (query.useCache !== false) {
        this.cacheQuery(cacheKey, context);
      }

      return context;

    } catch (error) {
      throw new Error(`Memory retrieval failed: ${error instanceof Error ? error.message : 'Unknown error'}`);
    }
  }

  // ==================== SPECIFIC SEARCH IMPLEMENTATIONS ====================

  /**
   * Vector similarity search using HNSW index
   */
  private async performVectorSearch(
    query: UnifiedMemoryQuery, 
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    const vectorStart = Date.now();
    
    if (!query.query) {
      throw new Error('Vector search requires a query string');
    }

    // Generate query embedding
    const embedding = await this.embeddingService.embedText({
      text: query.query,
      taskType: 'RETRIEVAL_QUERY'
    });

    // Search vector index
    const vectorResults = await this.vectorManager.searchSimilarTexts(
      'default-user',
      query.query,
      {
        k: query.similarity?.k || 10,
        efSearch: query.similarity?.efSearch || 100,
        threshold: query.similarity?.threshold || 0.7
      }
    );

    stats.vectorSearchTime = Date.now() - vectorStart;

    // Convert to unified format
    return await this.convertVectorResults(vectorResults.results, query);
  }

  /**
   * Semantic search with AI understanding
   */
  private async performSemanticSearch(
    query: UnifiedMemoryQuery,
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    const semanticStart = Date.now();

    // First, get vector results as base
    const vectorResults = await this.performVectorSearch(query, stats);

    // Apply semantic analysis for better understanding
    if (query.query) {
      // Analyze query semantics
      const semanticAnalysis = await this.analyzeQuerySemantics(query.query);
      
      // Re-rank results based on semantic understanding
      for (const result of vectorResults) {
        result.searchRelevance.semanticMatch = await this.calculateSemanticMatch(
          result, 
          semanticAnalysis
        );
        result.relevanceScore = this.combineRelevanceScores(result.searchRelevance);
      }
    }

    stats.semanticAnalysisTime = (stats.semanticAnalysisTime || 0) + (Date.now() - semanticStart);
    
    return vectorResults;
  }

  /**
   * Knowledge graph traversal search
   */
  private async performGraphSearch(
    query: UnifiedMemoryQuery,
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    const graphStart = Date.now();

    const results: UnifiedMemoryResult[] = [];

    if (query.query) {
      // Search graph for related entities and memories
      const graphResults = await this.graphManager.searchGraph(
        query.query,
        {
          maxResults: query.similarity?.k || 20
        }
      );

      // Find memories related to discovered entities
      for (const entity of graphResults.entities) {
        const relatedMemories = await this.graphManager.findMemoriesRelatedToEntity(
          entity.label || entity.id,
          query.userId
        );
        
        const memoryIds = Array.isArray(relatedMemories) ? relatedMemories : 
          (relatedMemories?.memories || []);
        for (const memoryId of memoryIds) {
          const memory = await this.retrieveMemoryById(memoryId, query);
          if (memory) {
            memory.searchRelevance.graphRelevance = entity.confidence || 0.8;
            results.push(memory);
          }
        }
      }
    }

    stats.graphTraversalTime = Date.now() - graphStart;
    return results;
  }

  /**
   * Temporal/time-based search
   */
  private async performTemporalSearch(
    query: UnifiedMemoryQuery,
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    // Get all user memories
    const allMemories = await this.getAllUserMemories(query.userId);
    
    // Filter by date range
    let filteredMemories = allMemories;
    
    if (query.dateRange?.start) {
      filteredMemories = filteredMemories.filter(
        m => m.created >= query.dateRange!.start!
      );
    }
    
    if (query.dateRange?.end) {
      filteredMemories = filteredMemories.filter(
        m => m.created <= query.dateRange!.end!
      );
    }

    // Calculate temporal relevance
    filteredMemories.forEach(memory => {
      memory.searchRelevance.temporalRelevance = this.calculateTemporalRelevance(
        memory.created,
        query.dateRange
      );
    });

    return filteredMemories;
  }

  /**
   * Keyword-based search
   */
  private async performKeywordSearch(
    query: UnifiedMemoryQuery,
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    if (!query.query) {
      return [];
    }

    const keywords = query.query.toLowerCase().split(/\s+/);
    const allMemories = await this.getAllUserMemories(query.userId);
    
    // Score memories based on keyword matches
    const results = allMemories.filter(memory => {
      const content = (memory.content || '').toLowerCase();
      const matchCount = keywords.filter(keyword => 
        content.includes(keyword)
      ).length;
      
      if (matchCount > 0) {
        memory.searchRelevance.keywordMatch = matchCount / keywords.length;
        return true;
      }
      
      return false;
    });

    return results;
  }

  /**
   * Hybrid search combining multiple methods
   */
  private async performHybridSearch(
    query: UnifiedMemoryQuery,
    stats: Partial<RetrievalStats['searchBreakdown']>
  ): Promise<UnifiedMemoryResult[]> {
    const allResults: UnifiedMemoryResult[] = [];
    const resultMap = new Map<string, UnifiedMemoryResult>();

    // Perform multiple search types
    if (query.query) {
      // Vector search
      const vectorResults = await this.performVectorSearch(query, stats);
      vectorResults.forEach(result => {
        result.searchRelevance.vectorSimilarity = result.similarity || 0;
        resultMap.set(result.id, result);
      });

      // Semantic enhancement
      const semanticResults = await this.performSemanticSearch(query, stats);
      semanticResults.forEach(result => {
        const existing = resultMap.get(result.id);
        if (existing) {
          existing.searchRelevance.semanticMatch = result.searchRelevance.semanticMatch;
        } else {
          resultMap.set(result.id, result);
        }
      });

      // Graph search
      const graphResults = await this.performGraphSearch(query, stats);
      graphResults.forEach(result => {
        const existing = resultMap.get(result.id);
        if (existing) {
          existing.searchRelevance.graphRelevance = result.searchRelevance.graphRelevance;
        } else {
          resultMap.set(result.id, result);
        }
      });
    }

    // Temporal search if date range specified
    if (query.dateRange) {
      const temporalResults = await this.performTemporalSearch(query, stats);
      temporalResults.forEach(result => {
        const existing = resultMap.get(result.id);
        if (existing) {
          existing.searchRelevance.temporalRelevance = result.searchRelevance.temporalRelevance;
        } else {
          resultMap.set(result.id, result);
        }
      });
    }

    // Combine and calculate final relevance scores
    Array.from(resultMap.values()).forEach(result => {
      result.relevanceScore = this.combineRelevanceScores(result.searchRelevance);
    });

    return Array.from(resultMap.values());
  }

  // ==================== HELPER METHODS ====================

  /**
   * Apply filters to search results
   */
  private async applyFilters(
    results: UnifiedMemoryResult[],
    query: UnifiedMemoryQuery
  ): Promise<UnifiedMemoryResult[]> {
    let filtered = results;

    // Category filter
    if (query.categories && query.categories.length > 0) {
      filtered = filtered.filter(result => 
        query.categories!.includes(result.category)
      );
    }

    // Importance filter
    if (query.importanceRange) {
      filtered = filtered.filter(result => {
        const importance = result.metadata.importance;
        return (!query.importanceRange!.min || importance >= query.importanceRange!.min) &&
               (!query.importanceRange!.max || importance <= query.importanceRange!.max);
      });
    }

    // Tags filter
    if (query.tags && query.tags.length > 0) {
      filtered = filtered.filter(result =>
        query.tags!.some(tag => result.metadata.tags.includes(tag))
      );
    }

    // Content type filter
    if (query.contentTypes && query.contentTypes.length > 0) {
      filtered = filtered.filter(result =>
        query.contentTypes!.includes(result.metadata.contentType)
      );
    }

    return filtered;
  }

  /**
   * Enrich results with additional data
   */
  private async enrichResults(
    results: UnifiedMemoryResult[],
    query: UnifiedMemoryQuery
  ): Promise<UnifiedMemoryResult[]> {
    const enriched = await Promise.all(results.map(async (result) => {
      // Load content if requested
      if (query.includeContent && !result.content) {
        try {
          const content = await this.loadMemoryContent(result.storage.blobId, query.userId);
          result.content = content;
        } catch (error) {
          console.warn(`Failed to load content for memory ${result.id}:`, error);
        }
      }

      // Load analytics if requested
      if (query.includeAnalytics) {
        result.analytics = await this.loadMemoryAnalytics(result.id);
      }

      // Load relationships if requested
      if (query.includeMetadata) {
        result.relationships = await this.loadMemoryRelationships(result.id, query.userId);
      }

      return result;
    }));

    return enriched;
  }

  // ... [Continuing with remaining helper methods]

  /**
   * Rank and sort results by relevance
   */
  private rankAndSortResults(
    results: UnifiedMemoryResult[],
    query: UnifiedMemoryQuery
  ): UnifiedMemoryResult[] {
    return results
      .sort((a, b) => b.relevanceScore - a.relevanceScore)
      .slice(0, query.similarity?.k || 50);
  }

  /**
   * Generate comprehensive search statistics
   */
  private generateStats(
    results: UnifiedMemoryResult[],
    searchBreakdown: Partial<RetrievalStats['searchBreakdown']>,
    startTime: number
  ): RetrievalStats {
    const totalTime = Date.now() - startTime;
    
    return {
      totalResults: results.length,
      processingTime: totalTime,
      cacheHitRate: this.calculateCacheHitRate(),
      searchBreakdown: {
        vectorSearchTime: searchBreakdown.vectorSearchTime || 0,
        semanticAnalysisTime: searchBreakdown.semanticAnalysisTime || 0,
        decryptionTime: searchBreakdown.decryptionTime || 0,
        graphTraversalTime: searchBreakdown.graphTraversalTime || 0
      },
      dataSourcesUsed: this.getDataSourcesUsed(results),
      qualityMetrics: {
        averageRelevance: results.reduce((sum, r) => sum + r.relevanceScore, 0) / results.length || 0,
        diversityScore: this.calculateDiversityScore(results),
        freshnessScore: this.calculateFreshnessScore(results)
      }
    };
  }

  // ==================== CACHE MANAGEMENT ====================

  private generateCacheKey(query: UnifiedMemoryQuery): string {
    return `query:${JSON.stringify(query)}`;
  }

  private getCachedQuery(key: string): RetrievalContext | null {
    const cached = this.queryCache.get(key);
    if (cached && Date.now() - cached.timestamp < this.QUERY_CACHE_TTL) {
      return cached.result;
    }
    this.queryCache.delete(key);
    return null;
  }

  private cacheQuery(key: string, result: RetrievalContext): void {
    this.queryCache.set(key, { result, timestamp: Date.now() });
  }

  // ==================== UTILITY METHODS ====================

  private calculateCacheHitRate(): number {
    // Implementation for cache hit rate calculation
    return 0.85; // Placeholder
  }

  private getDataSourcesUsed(results: UnifiedMemoryResult[]): string[] {
    const sources = new Set<string>();
    results.forEach(result => {
      if (result.storage.cacheStatus === 'cached') sources.add('cache');
      if (result.storage.walrusHash) sources.add('walrus');
      if (result.storage.blockchainId) sources.add('blockchain');
    });
    return Array.from(sources);
  }

  private calculateDiversityScore(results: UnifiedMemoryResult[]): number {
    const categories = new Set(results.map(r => r.category));
    return categories.size / Math.max(results.length, 1);
  }

  private calculateFreshnessScore(results: UnifiedMemoryResult[]): number {
    const now = Date.now();
    const avgAge = results.reduce((sum, r) => sum + (now - r.created.getTime()), 0) / results.length;
    const maxAge = 365 * 24 * 60 * 60 * 1000; // 1 year in ms
    return Math.max(0, 1 - (avgAge / maxAge));
  }

  // Placeholder methods - to be implemented
  private async convertVectorResults(vectorResults: any[], query: UnifiedMemoryQuery): Promise<UnifiedMemoryResult[]> {
    // Convert vector search results to unified format
    return [];
  }

  private async analyzeQuerySemantics(query: string): Promise<any> {
    // Analyze query for semantic understanding
    return {};
  }

  private async calculateSemanticMatch(result: UnifiedMemoryResult, semantics: any): Promise<number> {
    // Calculate semantic match score
    return 0.8;
  }

  private combineRelevanceScores(scores: UnifiedMemoryResult['searchRelevance']): number {
    // Combine different relevance scores into final score
    const weights = {
      vectorSimilarity: 0.4,
      semanticMatch: 0.3,
      keywordMatch: 0.1,
      graphRelevance: 0.15,
      temporalRelevance: 0.05
    };

    return (scores.vectorSimilarity || 0) * weights.vectorSimilarity +
           (scores.semanticMatch || 0) * weights.semanticMatch +
           (scores.keywordMatch || 0) * weights.keywordMatch +
           (scores.graphRelevance || 0) * weights.graphRelevance +
           (scores.temporalRelevance || 0) * weights.temporalRelevance;
  }

  private calculateTemporalRelevance(date: Date, range?: { start?: Date; end?: Date }): number {
    // Calculate temporal relevance score
    return 0.7;
  }

  private async getAllUserMemories(userId: string): Promise<UnifiedMemoryResult[]> {
    // Get all memories for user from blockchain
    return [];
  }

  private async retrieveMemoryById(memoryId: string, query: UnifiedMemoryQuery): Promise<UnifiedMemoryResult | null> {
    // Retrieve specific memory by ID
    return null;
  }

  private async loadMemoryContent(blobId: string, userId: string): Promise<string> {
    // Load and decrypt memory content
    return "";
  }

  private async loadMemoryAnalytics(memoryId: string): Promise<any> {
    // Load memory analytics data
    return {};
  }

  private async loadMemoryRelationships(memoryId: string, userId: string): Promise<any> {
    // Load memory relationships from knowledge graph
    return {};
  }

  private async generateSuggestions(query: UnifiedMemoryQuery, results: UnifiedMemoryResult[]): Promise<any> {
    // Generate query suggestions and exploration paths
    return {
      relatedQueries: [],
      filterSuggestions: [],
      explorationPaths: []
    };
  }

  private async generateTimeline(results: UnifiedMemoryResult[], query: UnifiedMemoryQuery): Promise<any> {
    // Generate timeline of memory events
    return undefined;
  }

  // ==================== PUBLIC API ====================

  /**
   * Get comprehensive memory retrieval statistics
   */
  getRetrievalStats(): {
    cacheStats: { queries: number; content: number; analytics: number };
    performanceMetrics: any;
  } {
    return {
      cacheStats: {
        queries: this.queryCache.size,
        content: this.contentCache.size,
        analytics: this.analyticsCache.size
      },
      performanceMetrics: {}
    };
  }

  /**
   * Clear all caches
   */
  clearCache(): void {
    this.queryCache.clear();
    this.contentCache.clear();
    this.analyticsCache.clear();
  }

  /**
   * Warm up caches for a user
   */
  async warmupCache(userId: string): Promise<void> {
    // Pre-load frequently accessed data
    console.log(`Warming up cache for user ${userId}`);
  }
}