UNPKG

@escher-dbai/rag-module

Version:

Enterprise RAG module with chat context storage, vector search, and session management. Complete chat history retrieval and streaming content extraction for Electron apps.

818 lines (700 loc) 26.2 kB
/** * SearchService - Collection-specific search per business architecture * Chat History: Scroll API for filter-only search (1D dummy vectors) * AWS Estate: Vector search for semantic + filters (1024D BGE-M3 vectors) */ class SearchService { constructor(vectorStore, embeddingService, securityService, mappingService = null, iamService = null) { this.vectorStore = vectorStore; this.embeddingService = embeddingService; this.securityService = securityService; this.mappingService = mappingService; this.iamService = iamService; // Collection-specific configurations per business specs this.collectionConfigs = { chat: { useVectorSearch: false, // Scroll API only for chat history requiresEmbedding: false, defaultLimit: 50, maxLimit: 200, defaultFilters: ['context_id', 'role', 'timestamp'] }, estate: { useVectorSearch: true, // Vector search for estate resources requiresEmbedding: true, defaultLimit: 15, maxLimit: 50, defaultFilters: ['resource_type', 'account_id', 'region', 'service'] } }; } // ============ BUSINESS ARCHITECTURE SEARCH METHODS ============ /** * Search chat history using scroll API (Business Architecture API) * @param {Object} options - Chat search options * @param {string} options.contextId - Conversation context ID * @param {string} [options.role] - Message role filter * @param {number} [options.fromTimestamp] - Start timestamp * @param {number} [options.toTimestamp] - End timestamp * @param {number} [options.limit] - Result limit * @returns {Promise<Array>} Chat message results */ async searchChatHistory(options = {}) { const config = this.collectionConfigs.chat; const limit = Math.min(options.limit || config.defaultLimit, config.maxLimit); // Build chat-specific filters const filters = this._buildChatFilters(options); try { // Use scroll API for chat (no vector search needed) const results = await this.vectorStore.scrollCollection('chat', { filter: filters, limit, withPayload: true, withVector: false // Don't need dummy vectors in results }); // Process and decrypt chat messages return this._processChatResults(results, options); } catch (error) { console.error('❌ Chat history search failed:', error.message); throw new Error(`Chat search failed: ${error.message}`); } } /** * Search AWS estate using vector search (Business Architecture API) * @param {string} query - Semantic search query * @param {Object} options - Estate search options * @param {string[]} [options.resourceTypes] - Resource type filters * @param {string[]} [options.accountIds] - Account ID filters * @param {string[]} [options.regions] - Region filters * @param {string[]} [options.services] - Service filters * @param {Object} [options.iamContext] - IAM context for permissions * @param {number} [options.limit] - Result limit * @returns {Promise<Array>} Estate resource results */ async searchEstateResources(query, options = {}) { const config = this.collectionConfigs.estate; const limit = Math.min(options.limit || config.defaultLimit, config.maxLimit); // Preprocess query for estate search const preprocessedQuery = this._preprocessEstateQuery(query); // Generate query embedding for semantic search const queryEmbedding = await this.embeddingService.embed(preprocessedQuery); // Build estate-specific filters const filters = this._buildEstateFilters(options); try { // Use vector search for estate resources const searchLimit = this._calculateOptimalLimit(limit); const results = await this.vectorStore.searchCollection('estate', queryEmbedding, { limit: searchLimit, scoreThreshold: options.scoreThreshold || 0.2, filter: filters, withPayload: true, withVector: false // Don't need vectors in results }); // Apply relevance filtering const filteredResults = this._applyRelevanceFiltering(results, limit); // Apply IAM permission filtering if context provided const permissionFilteredResults = await this._applyIAMFiltering(filteredResults, options.iamContext); // Process and decrypt estate resources return this._processEstateResults(permissionFilteredResults, options); } catch (error) { console.error('❌ Estate search failed:', error.message); throw new Error(`Estate search failed: ${error.message}`); } } /** * Legacy search method - routes to appropriate collection * @deprecated Use searchChatHistory or searchEstateResources instead */ async search(query, options = {}) { console.warn('⚠️ Using legacy search method. Use searchChatHistory or searchEstateResources instead.'); // Route based on options or query characteristics if (options.collectionType === 'chat' || options.contextId) { return this.searchChatHistory(options); } else { return this.searchEstateResources(query, options); } } // ============ COLLECTION-SPECIFIC FILTER BUILDERS ============ /** * Build chat history filters per business architecture * @param {Object} options - Chat search options * @returns {Object} Qdrant filter structure */ _buildChatFilters(options) { const mustConditions = []; // Required context_id filter for chat history if (options.contextId) { mustConditions.push({ key: 'context_id', match: { value: options.contextId } }); } // Optional role filter if (options.role) { mustConditions.push({ key: 'role', match: { value: options.role } }); } // Timestamp range filter if (options.fromTimestamp || options.toTimestamp) { const rangeFilter = { key: 'timestamp', range: {} }; if (options.fromTimestamp) rangeFilter.range.gte = options.fromTimestamp; if (options.toTimestamp) rangeFilter.range.lte = options.toTimestamp; mustConditions.push(rangeFilter); } // Message index range (for pagination) if (options.fromMessageIndex !== undefined || options.toMessageIndex !== undefined) { const indexFilter = { key: 'message_index', range: {} }; if (options.fromMessageIndex !== undefined) indexFilter.range.gte = options.fromMessageIndex; if (options.toMessageIndex !== undefined) indexFilter.range.lte = options.toMessageIndex; mustConditions.push(indexFilter); } return mustConditions.length > 0 ? { must: mustConditions } : {}; } /** * Build estate resource filters per business architecture * @param {Object} options - Estate search options * @returns {Object} Qdrant filter structure */ _buildEstateFilters(options) { const mustConditions = []; // Resource type filter if (options.resourceTypes && options.resourceTypes.length > 0) { if (options.resourceTypes.length === 1) { mustConditions.push({ key: 'resource_type', match: { value: options.resourceTypes[0] } }); } else { mustConditions.push({ key: 'resource_type', match: { any: options.resourceTypes } }); } } // Account ID filter if (options.accountIds && options.accountIds.length > 0) { if (options.accountIds.length === 1) { mustConditions.push({ key: 'account_id', match: { value: options.accountIds[0] } }); } else { mustConditions.push({ key: 'account_id', match: { any: options.accountIds } }); } } // Region filter if (options.regions && options.regions.length > 0) { if (options.regions.length === 1) { mustConditions.push({ key: 'region', match: { value: options.regions[0] } }); } else { mustConditions.push({ key: 'region', match: { any: options.regions } }); } } // Service filter if (options.services && options.services.length > 0) { if (options.services.length === 1) { mustConditions.push({ key: 'service', match: { value: options.services[0] } }); } else { mustConditions.push({ key: 'service', match: { any: options.services } }); } } // State filter if (options.states && options.states.length > 0) { if (options.states.length === 1) { mustConditions.push({ key: 'state', match: { value: options.states[0] } }); } else { mustConditions.push({ key: 'state', match: { any: options.states } }); } } // Tag filters (selective indexing per business specs) if (options.environment) { mustConditions.push({ key: 'tags.env', match: { value: options.environment } }); } if (options.application) { mustConditions.push({ key: 'tags.app', match: { value: options.application } }); } // Last synced filter (for freshness) if (options.syncedAfter) { mustConditions.push({ key: 'last_synced', range: { gte: options.syncedAfter } }); } return mustConditions.length > 0 ? { must: mustConditions } : {}; } /** * Build estate resource filters per business architecture * @param {Object} options - Estate search options * @returns {Object} Qdrant filter structure */ _buildEstateFilters(options) { const mustConditions = []; // Resource type filter if (options.resourceTypes && options.resourceTypes.length > 0) { if (options.resourceTypes.length === 1) { mustConditions.push({ key: 'resource_type', match: { value: options.resourceTypes[0] } }); } else { mustConditions.push({ key: 'resource_type', match: { any: options.resourceTypes } }); } } // Account ID filter if (options.accountIds && options.accountIds.length > 0) { if (options.accountIds.length === 1) { mustConditions.push({ key: 'account_id', match: { value: options.accountIds[0] } }); } else { mustConditions.push({ key: 'account_id', match: { any: options.accountIds } }); } } // Region filter if (options.regions && options.regions.length > 0) { if (options.regions.length === 1) { mustConditions.push({ key: 'region', match: { value: options.regions[0] } }); } else { mustConditions.push({ key: 'region', match: { any: options.regions } }); } } // Service filter if (options.services && options.services.length > 0) { if (options.services.length === 1) { mustConditions.push({ key: 'service', match: { value: options.services[0] } }); } else { mustConditions.push({ key: 'service', match: { any: options.services } }); } } // State filter if (options.states && options.states.length > 0) { if (options.states.length === 1) { mustConditions.push({ key: 'state', match: { value: options.states[0] } }); } else { mustConditions.push({ key: 'state', match: { any: options.states } }); } } // Tag filters (selective indexing per business specs) if (options.environment) { mustConditions.push({ key: 'tags.env', match: { value: options.environment } }); } if (options.application) { mustConditions.push({ key: 'tags.app', match: { value: options.application } }); } // Last synced filter (for freshness) if (options.syncedAfter) { mustConditions.push({ key: 'last_synced', range: { gte: options.syncedAfter } }); } return mustConditions.length > 0 ? { must: mustConditions } : {}; } /** * Legacy filter builder for backward compatibility * @deprecated Use _buildChatFilters or _buildEstateFilters */ _buildFilter(options) { console.warn('⚠️ Using legacy _buildFilter. Use collection-specific filter builders.'); // Simple conversion for backward compatibility const filter = {}; if (options.clouds && Array.isArray(options.clouds)) { filter.cloud = options.clouds.length === 1 ? options.clouds[0] : options.clouds; } if (options.services && Array.isArray(options.services)) { filter.service = options.services.length === 1 ? options.services[0] : options.services; } if (options.regions && Array.isArray(options.regions)) { filter.region = options.regions.length === 1 ? options.regions[0] : options.regions; } if (options.environments && Array.isArray(options.environments)) { filter.environment = options.environments.length === 1 ? options.environments[0] : options.environments; } if (options.filters && typeof options.filters === 'object') { Object.assign(filter, options.filters); } return filter; } // ============ QUERY PREPROCESSING ============ /** * Preprocess estate search queries with AWS/Azure/GCP spell correction * @param {string} query - Raw search query * @returns {string} Preprocessed query */ _preprocessEstateQuery(query) { if (!query) return ''; query = query.toLowerCase().trim(); // Extended spell corrections for multi-cloud estate const spellCorrections = { // AWS corrections 'lamda': 'lambda', 'lambada': 'lambda', 'lamabda': 'lambda', 'lamnda': 'lambda', 'dynamo': 'dynamodb', 'api-gateway': 'api gateway', 'apigateway': 'api gateway', 'secrets': 'secrets manager', 'parameter': 'parameter store', 'systems': 'systems manager', 'certificate': 'certificate manager', 'directory': 'directory service', 'control-tower': 'control tower', 'security-hub': 'security hub', 'disaster': 'disaster recovery', 'running instances': 'running instance ec2', 'stopped instances': 'stopped instance ec2', 'database instances': 'database instance rds', 'storage buckets': 'storage bucket s3', // Azure corrections 'virtual machines': 'virtual machine vm', 'storage accounts': 'storage account', 'app services': 'app service', 'key vaults': 'key vault', 'sql databases': 'sql database', 'service bus': 'servicebus', 'cosmos db': 'cosmosdb', 'azure ad': 'active directory', // GCP corrections 'compute engine': 'compute-engine vm', 'cloud storage': 'cloud-storage bucket', 'cloud functions': 'cloud-functions', 'cloud sql': 'cloud-sql database', 'pub sub': 'pub-sub', 'cloud run': 'cloud-run', 'kubernetes engine': 'gke cluster' }; // Apply word-level corrections const words = query.split(/\s+/); for (let i = 0; i < words.length; i++) { const word = words[i].replace(/[.,!?;:]/g, ''); if (spellCorrections[word]) { words[i] = spellCorrections[word]; } } // Apply phrase-level corrections let correctedQuery = words.join(' '); for (const [mistake, correction] of Object.entries(spellCorrections)) { if (correctedQuery.includes(mistake)) { correctedQuery = correctedQuery.replace(new RegExp(mistake, 'gi'), correction); } } return correctedQuery; } /** * Legacy preprocessor for backward compatibility * @deprecated Use _preprocessEstateQuery */ _preprocessQuery(query) { console.warn('⚠️ Using legacy _preprocessQuery. Use _preprocessEstateQuery for estate searches.'); return this._preprocessEstateQuery(query); } /** * Calculate optimal search limit (from traditional RAG) */ _calculateOptimalLimit(requestedLimit) { // Normalize the limit (minimum 10, maximum 50 for performance) const normalizedLimit = Math.max(10, Math.min(requestedLimit, 50)); // Apply over-fetch for filtering (like traditional RAG balanced mode) return normalizedLimit * 5; // Even more aggressive over-fetch for better recall } /** * Apply relevance filtering (from traditional RAG) with deduplication */ _applyRelevanceFiltering(results, requestedLimit) { if (results.length === 0) { return results; } // Constants from traditional RAG - more permissive for better recall const DYNAMIC_THRESHOLD_RATIO = 0.75; // Keep results ≥ 75% of best score (was 85%) const ABSOLUTE_MIN_SCORE = 0.20; // Absolute minimum score const MAX_RESULTS = Math.min(requestedLimit, 20); // Maximum results const SCORE_DROP_THRESHOLD = 0.35; // Allow larger score drops (was 25%) const filtered = []; const seenIds = new Set(); // Deduplication set const bestScore = results[0]?.score || 0; // Calculate dynamic threshold: 75% of best score, but at least 0.20 let dynamicThreshold = bestScore * DYNAMIC_THRESHOLD_RATIO; if (dynamicThreshold < ABSOLUTE_MIN_SCORE) { dynamicThreshold = ABSOLUTE_MIN_SCORE; } for (let i = 0; i < results.length && filtered.length < MAX_RESULTS; i++) { const result = results[i]; // Skip duplicates by document ID if (seenIds.has(result.id)) { continue; } // Apply dynamic relevance threshold if (result.score < dynamicThreshold) { continue; } // Apply score drop threshold if (bestScore - result.score > SCORE_DROP_THRESHOLD) { continue; } // Add to results and mark as seen seenIds.add(result.id); filtered.push(result); } return filtered; } // ============ COLLECTION-SPECIFIC RESULT PROCESSORS ============ /** * Process chat history search results * @param {Array} results - Raw chat results from scroll API * @param {Object} options - Processing options * @returns {Promise<Array>} Processed chat messages */ async _processChatResults(results, options = {}) { if (!results || results.length === 0) { return []; } const processedResults = []; for (const result of results) { try { const processedResult = { id: result.id, context_id: result.payload.context_id || result.payload.contextId, message_index: result.payload.message_index, role: result.payload.role, timestamp: result.payload.timestamp, // Content should already be decrypted by the embedded backend // For server backends, it might be in encrypted_content field content: result.payload.content || result.payload.encrypted_content, // Include metadata if requested ...(options.includeMetadata && { metadata: { created_at: new Date(result.payload.timestamp), collection_type: 'chat' } }) }; processedResults.push(processedResult); } catch (error) { console.error(`❌ Failed to process chat result ${result.id}:`, error.message); // Skip failed results rather than failing the entire search continue; } } // Sort by message_index for proper conversation order processedResults.sort((a, b) => a.message_index - b.message_index); return processedResults; } /** * Process estate resource search results * @param {Array} results - Raw estate results from vector search * @param {Object} options - Processing options * @returns {Promise<Array>} Processed estate resources */ async _processEstateResults(results, options = {}) { if (!results || results.length === 0) { return []; } const processedResults = []; for (const result of results) { try { // For embedded backend, data should already be decrypted // For server backend, we might need to handle encrypted_data let resourceData; if (result.payload.encrypted_data) { // This would need proper decryption service, but for now skip encryption resourceData = JSON.parse(result.payload.encrypted_data); } else if (result.payload.data) { resourceData = result.payload.data; } else { resourceData = {}; // Fallback for missing data } const processedResult = { id: result.id, score: result.score, // Semantic similarity score // Plain text fields (indexed for filtering) resource_type: result.payload.resource_type, account_id: result.payload.account_id, account_name: result.payload.account_name, region: result.payload.region, service: result.payload.service, state: result.payload.state, last_synced: result.payload.last_synced, // Tags (if present) tags: { ...(result.payload['tags.env'] && { env: result.payload['tags.env'] }), ...(result.payload['tags.app'] && { app: result.payload['tags.app'] }), ...(result.payload['tags.name'] && { name: result.payload['tags.name'] }) }, // Decrypted sensitive data identifier: resourceData.identifier, arn: resourceData.arn, name: resourceData.name, // IAM permissions if available ...(resourceData.iam && { permissions: resourceData.iam }), // Include full metadata if requested ...(options.includeMetadata && { metadata: { ...resourceData.metadata, collection_type: 'estate', search_score: result.score } }) }; // Apply anonymous ID mapping if enabled if (this.mappingService && options.useAnonymousIds !== false) { processedResult.anonymous_id = await this.mappingService.getAnonymousId(result.id); } processedResults.push(processedResult); } catch (error) { console.error(`❌ Failed to process estate result ${result.id}:`, error.message); // Skip failed results rather than failing the entire search continue; } } // Results are already sorted by score from vector search return processedResults; } /** * Apply IAM permission filtering to estate results * @param {Array} results - Estate search results * @param {Object} iamContext - IAM context (user, roles, permissions) * @returns {Promise<Array>} Permission-filtered results */ async _applyIAMFiltering(results, iamContext) { if (!this.iamService || !iamContext || !results.length) { return results; // No IAM filtering if service not available or no context } try { const filteredResults = []; for (const result of results) { // Check if user has permission to view this resource const hasAccess = await this.iamService.checkResourceAccess( iamContext, result.id, 'read' ); if (hasAccess) { filteredResults.push(result); } } return filteredResults; } catch (error) { console.error('❌ IAM filtering failed:', error.message); // On IAM errors, return original results (fail open for availability) console.warn('⚠️ Returning unfiltered results due to IAM error'); return results; } } /** * Legacy result processor for backward compatibility * @deprecated Use _processChatResults or _processEstateResults */ async _processResults(results, options) { console.warn('⚠️ Using legacy _processResults. Use collection-specific processors.'); // Apply anonymous ID mapping if enabled if (this.mappingService && options.useAnonymousIds !== false) { results = await this._applyAnonymousMapping(results); } if (options.externalFormat) { return this.securityService.formatForExternal(results); } return results; } /** * Apply anonymous ID mapping to search results */ async _applyAnonymousMapping(results) { const mappedResults = []; for (const result of results) { const mappedResult = { ...result }; // Map the main ID if it looks like an ARN if (result.id && result.id.startsWith('arn:aws:')) { const anonymousId = await this.mappingService.getAnonymousId(result.id); if (anonymousId) { mappedResult.id = anonymousId; mappedResult.original_id = result.id; // Keep for internal use } } // Map any ARNs in metadata if (result.metadata) { mappedResult.metadata = await this._mapMetadataArns(result.metadata); } mappedResults.push(mappedResult); } return mappedResults; } /** * Map ARNs in metadata to anonymous IDs */ async _mapMetadataArns(metadata) { const mappedMetadata = { ...metadata }; for (const [key, value] of Object.entries(mappedMetadata)) { if (typeof value === 'string' && value.startsWith('arn:aws:')) { const anonymousId = await this.mappingService.getAnonymousId(value); if (anonymousId) { mappedMetadata[key] = anonymousId; } } } return mappedMetadata; } } module.exports = SearchService;