RAG_enhancer.py

import torch
import torch.nn as nn
import numpy as np
from typing import List, Dict, Set, Optional, Tuple, Callable
import scipy.sparse as sp
from tqdm import tqdm
from torch_geometric.data import Data
from torch_geometric.utils import k_hop_subgraph
import os
import json
import hashlib
import logging
from dataclasses import dataclass
from search_nano_db import RAGRetriever
from feature_enhancer import FeatureEnhancer
from search_motif_db import MotifRetriever
from centrality_utils import compute_centrality_and_cse

LOWERCASE_DATASETS = {'wikics', 'home'}

@dataclass
class RAGConfig:
    k: int = 5
    dataset_name: str = "default"
    cache_dir: str = "motif_lib"
    enable_cse_cache: bool = True
    cse_ksteps: List[int] = None
    feature_dim: int = 100 
    lowercase_datasets: Set[str] = None
    
    def __post_init__(self):
        if self.cse_ksteps is None:
            self.cse_ksteps = list(range(1, 9))
        if self.lowercase_datasets is None:
            self.lowercase_datasets = LOWERCASE_DATASETS

class RAGEnhancer:
    def __init__(self, config: Optional[RAGConfig] = None, feature_dim: Optional[int] = None):

        self.config = config or RAGConfig()
        if feature_dim is not None:
            self.config.feature_dim = feature_dim
        self.logger = self._setup_logger()
        
        self.rag_retriever = None
        self.feature_enhancer = None
        self.motif_retriever = None
        
    
    def _setup_logger(self) -> logging.Logger:
        logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger
    
    def initialize_components(self, 
                            pretrain_dataset_names: Set[str],
                            pretrain_weights: Optional[Dict[str, torch.Tensor]] = None,
                            db_file_path: str = "nano_db/unified_database.json",
                            text_encoder: str = "bert",
                            motif_lib_path: str = "motif_lib",
                            motif_db_path: str = "motif_lib") -> None:

        try:
            self.rag_retriever = RAGRetriever(
                db_file_path=db_file_path,
                text_encoder=text_encoder,
                device="cuda" if torch.cuda.is_available() else "cpu"
            )
            
            self.feature_enhancer = FeatureEnhancer(
                text_encoder=text_encoder,
                text_output_dim=50,
                device="cuda" if torch.cuda.is_available() else "cpu"
            )
            
            self.motif_retriever = MotifRetriever(
                motif_lib_path=motif_lib_path,
                motif_db_path=motif_db_path,
                device="cuda" if torch.cuda.is_available() else "cpu"
            )
            
        except Exception as e:
            self.logger.error(f"component initialization failed: {e}")
            raise
    
    def load_graph_data(self, 
                       edge_index: torch.Tensor, 
                       num_nodes: int,
                       raw_texts: Optional[List[str]] = None) -> Dict[str, any]:

        try:
            
            adj = self._build_adjacency_matrix(edge_index, num_nodes)
            
            processed_texts = self._preprocess_texts(raw_texts) if raw_texts else None
            
            graph_data = {
                'edge_index': edge_index,
                'num_nodes': num_nodes,
                'adj': adj,
                'raw_texts': raw_texts,
                'processed_texts': processed_texts
            }
            
            return graph_data
            
        except Exception as e:
            self.logger.error(f"loading graph data failed: {e}")
            raise
    
    def _build_adjacency_matrix(self, edge_index: torch.Tensor, num_nodes: int) -> sp.csr_matrix:
        try:
            adj = sp.csr_matrix(
                (np.ones(edge_index.shape[1]), (edge_index[0], edge_index[1])),
                shape=(num_nodes, num_nodes)
            )
            
            adj = adj + adj.T
            adj.data = np.ones_like(adj.data)
            
            return adj
            
        except Exception as e:
            self.logger.error(f"building adjacency matrix failed: {e}")
            raise
    
    def _preprocess_texts(self, raw_texts: List[str]) -> List[str]:
        if not raw_texts:
            return []
        
        processed_texts = []
        for text in raw_texts:
            if text and isinstance(text, str):
                cleaned_text = text.strip()
                if cleaned_text:
                    processed_texts.append(cleaned_text)
                else:
                    processed_texts.append("")
            else:
                processed_texts.append("")
        
        return processed_texts
    
    def _get_cse_cache_path(self, dataset_name: str, edge_index: torch.Tensor, num_nodes: int) -> str:

        if not self.config.enable_cse_cache:
            return None
            
        edge_hash = hashlib.md5(edge_index.cpu().numpy().tobytes()).hexdigest()[:16]
        cache_dir = os.path.join(self.config.cache_dir, dataset_name)
        os.makedirs(cache_dir, exist_ok=True)
        return os.path.join(cache_dir, f"cse_features_{num_nodes}_{edge_hash}.pt")
    
    def _load_cse_features(self, cache_path: str) -> Optional[torch.Tensor]:

        if not cache_path or not self.config.enable_cse_cache:
            return None
            
        try:
            if os.path.exists(cache_path):
                cse_features = torch.load(cache_path, map_location='cpu')
                return cse_features
        except Exception as e:
            self.logger.warning(f"loading CSE cache failed: {e}")
        return None
    
    def _save_cse_features(self, cse_features: torch.Tensor, cache_path: str) -> None:

        if not cache_path or not self.config.enable_cse_cache:
            return
            
        try:
            torch.save(cse_features, cache_path)
        except Exception as e:
            self.logger.warning(f"saving CSE cache failed: {e}")

    def _get_neighbors(self, node_idx: int, adj: sp.csr_matrix) -> List[int]:

        if node_idx >= adj.shape[0]:
            return []
        return adj.indices[adj.indptr[node_idx]:adj.indptr[node_idx+1]].tolist()

    def _extract_1hop_subgraph(self, center_node_idx: int, edge_index: torch.Tensor, 
                              num_nodes: int, dataset_name: str = None) -> Data:

        if dataset_name is None:
            dataset_name = self.config.dataset_name
            
        subgraph_nodes, subgraph_edge_index, _, _ = k_hop_subgraph(
            node_idx=center_node_idx,
            num_hops=1,
            edge_index=edge_index,
            relabel_nodes=True,
            num_nodes=num_nodes
        )
        
        cache_path = self._get_cse_cache_path(dataset_name, edge_index, num_nodes)
        cse_features = self._load_cse_features(cache_path)
        
        if cse_features is None:
            cse_results = compute_centrality_and_cse(
                edge_index=edge_index,
                num_nodes=num_nodes,
                ksteps=self.config.cse_ksteps
            )
            cse_features = cse_results['cse_encodings']

            self._save_cse_features(cse_features, cache_path)
        
        subgraph_cse_features = cse_features[subgraph_nodes]
        
        subgraph = Data(
            x=subgraph_cse_features,
            edge_index=subgraph_edge_index,
            center_node_idx=center_node_idx,
            subgraph_nodes=subgraph_nodes,
            num_nodes=len(subgraph_nodes)
        )
        
        return subgraph
    

    @torch.no_grad()
    def enhance_features_for_shot_nodes(
        self,
        shot_node_indices: List[int],
        base_enhanced_features: torch.Tensor,
        raw_texts: List[str],
        adj: sp.csr_matrix,
        pretrain_dataset_names: Set[str],
        pretrain_weights: Optional[Dict[str, torch.Tensor]] = None,
        pretrain_features: Optional[Dict[str, torch.Tensor]] = None,
        edge_index: Optional[torch.Tensor] = None,
        k: int = None,
        dataset_name: str = None,
        rag_weight: float = 0.1,
        motif_weight: float = 0.1
    ) -> torch.Tensor:
        
        if k is None:
            k = self.config.k
        if dataset_name is None:
            dataset_name = self.config.dataset_name
            
        try:

            if not self._are_components_initialized():
                self.initialize_components(pretrain_dataset_names, pretrain_weights)
            
            if pretrain_features is not None:
                self.pretrain_features = pretrain_features
            
            nodes_to_enhance = set(shot_node_indices)
            for node_idx in shot_node_indices:
                neighbors = self._get_neighbors(node_idx, adj)
                nodes_to_enhance.update(neighbors)
            
            enhanced_features = self._step1_rag_enhancement(
                shot_node_indices=shot_node_indices,
                base_enhanced_features=base_enhanced_features,
                raw_texts=raw_texts,
                adj=adj,
                k=k,
                filter_datasets=pretrain_dataset_names,
                rag_weight=rag_weight
            )
            
            enhanced_features_motif = self._step2_motif_enhancement(
                shot_node_indices=shot_node_indices,
                enhanced_features=enhanced_features,
                adj=adj,
                pretrain_dataset_names=pretrain_dataset_names,
                pretrain_weights=pretrain_weights,
                edge_index=edge_index,
                dataset_name=dataset_name,
                motif_weight=motif_weight
            )
            
            return enhanced_features_motif
            
        except Exception as e:
            self.logger.error(f"feature enhancement failed: {e}")
            raise
    
    def _are_components_initialized(self) -> bool:

        return (self.rag_retriever is not None and 
                self.feature_enhancer is not None and
                self.motif_retriever is not None)
    
    def _normalize_dataset_name(self, dataset_name: str) -> str:

        if not dataset_name:
            return dataset_name
        
        dataset_lower = dataset_name.lower()
        
        if dataset_lower in self.config.lowercase_datasets:
            return dataset_lower
        else:
            return dataset_name.capitalize()
    
    def _step2_motif_enhancement(self,
                                shot_node_indices: List[int],
                                enhanced_features: torch.Tensor,
                                adj: sp.csr_matrix,
                                pretrain_dataset_names: Set[str],
                                pretrain_weights: Optional[Dict[str, torch.Tensor]] = None,
                                edge_index: Optional[torch.Tensor] = None,
                                dataset_name: str = None,
                                motif_weight: float = 0.15) -> torch.Tensor:

        if dataset_name is None:
            dataset_name = self.config.dataset_name
            
        motif_enhanced_features = enhanced_features.clone()
        
        nodes_to_enhance = set(shot_node_indices)
        for node_idx in shot_node_indices:
            neighbors = self._get_neighbors(node_idx, adj)
            nodes_to_enhance.update(neighbors)
        
        nodes_to_enhance = list(nodes_to_enhance)
        
        if edge_index is None:
            self.logger.warning("edge_index is not provided, cannot perform Motif enhancement")
            return motif_enhanced_features
        
        if pretrain_weights is None:
            self.logger.warning("pretrain weights are not provided, using default weights for Motif enhancement")
            pretrain_weights = {name: 1.0 for name in pretrain_dataset_names}
        
        for node_idx in nodes_to_enhance:
            try:
                motif_feature = self.get_multi_domain_motif_enhanced_feature(
                    node_idx=node_idx,
                    edge_index=edge_index,
                    num_nodes=enhanced_features.shape[0],
                    pretrain_dataset_names=pretrain_dataset_names,
                    pretrain_weights=pretrain_weights,
                    dataset_name=dataset_name,
                    motif_weight=motif_weight
                )
                
                if motif_feature is not None:
                    if motif_feature.device != motif_enhanced_features.device:
                        motif_feature = motif_feature.to(motif_enhanced_features.device)
                    
                    if motif_feature.shape[0] > self.config.feature_dim:
                        motif_feature = motif_feature[:self.config.feature_dim]
                    elif motif_feature.shape[0] < self.config.feature_dim:
                        padding = torch.zeros(self.config.feature_dim - motif_feature.shape[0], 
                                            device=motif_feature.device, 
                                            dtype=motif_feature.dtype)
                        motif_feature = torch.cat([motif_feature, padding], dim=0)
                    
                    motif_enhanced_features[node_idx][:self.config.feature_dim//2] += motif_weight * motif_feature[:self.config.feature_dim//2]
                    
                else:
                    # node not found Motif enhanced feature, skip
                    pass
                    
            except Exception as e:
                self.logger.error(f"node {node_idx} Motif enhancement failed: {e}")
                continue
        
        return motif_enhanced_features
    
    def get_related_feature_vectors(self, 
                                  node_id: int, 
                                  raw_texts: List[str],
                                  k: int = None,
                                  filter_datasets: Optional[Set[str]] = None) -> tuple[List[torch.Tensor], List[float]]:
        if k is None:
            k = self.config.k
            
        if node_id >= len(raw_texts) or node_id < 0:
            raise ValueError(f"node ID {node_id} out of range [0, {len(raw_texts)-1}]")
        
        if not self._are_components_initialized():
            raise RuntimeError("RAG components not initialized, please call initialize_components()")
        
        try:
            query_text = raw_texts[node_id]
            
            def dataset_filter(doc: Dict) -> bool:
                if filter_datasets is None:
                    return True
                
                metadata = doc.get("metadata", {})
                dataset = metadata.get("dataset")
                doc_type = metadata.get("type")
                
                if not dataset or not doc_type:
                    return False
                
                if doc_type != 'node':
                    return False
                
                dataset_lower = dataset.lower()
                filter_datasets_lower = {name.lower() for name in filter_datasets}
                
                
                return dataset_lower in filter_datasets_lower
            
            
            retrieved_results = self.rag_retriever.search(
                query_text=query_text,
                k=k,
                filter_lambda=dataset_filter if filter_datasets else None
            )
            
            if not retrieved_results:
                self.logger.warning("RAG retrieval returned empty results!")
            
            feature_vectors = []
            similarities = []
            
            if not retrieved_results:
                self.logger.warning("retrieval results are empty, cannot extract feature vectors")
                return [], []
                
            for i, result in enumerate(retrieved_results):
                metadata = result.get('metadata', {})
                dataset_name = self._normalize_dataset_name(metadata.get('dataset', ''))
                node_id_in_dataset = metadata.get('id_in_dataset')
                
                if dataset_name and node_id_in_dataset is not None:
                    pretrain_vector = self.get_pretrain_feature_vector(dataset_name, node_id_in_dataset)
                    
                    if pretrain_vector is not None:
                        feature_vectors.append(pretrain_vector)
                        raw_metrics = result.get('__metrics__', 1.0)
                        similarity = 1 - raw_metrics
                        similarities.append(similarity)
                    else:
                        self.logger.warning(f"  ✗ node {dataset_name}:{node_id_in_dataset} not found feature vector")
                else:
                    self.logger.warning(f"  ✗ skip invalid result - dataset: '{dataset_name}', node ID: {node_id_in_dataset}")
            
            if feature_vectors and similarities:
                sorted_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)
                
                sorted_feature_vectors = [feature_vectors[i] for i in sorted_indices]
                sorted_similarities = [similarities[i] for i in sorted_indices]
                
                return sorted_feature_vectors, sorted_similarities
            else:
                self.logger.warning("no valid feature vectors found")
                return [], []
            
        except Exception as e:
            self.logger.error(f"failed to get related feature vectors: {e}")
            return [], []
    
    def _step1_rag_enhancement(self, 
                              shot_node_indices: List[int],
                              base_enhanced_features: torch.Tensor,
                              raw_texts: List[str],
                              adj: sp.csr_matrix,
                              k: int,
                              filter_datasets: Set[str],
                              rag_weight: float = 0.125) -> torch.Tensor:
        
        enhanced_features = base_enhanced_features.clone()
        
        nodes_to_enhance = set(shot_node_indices)
        for node_idx in shot_node_indices:
            neighbors = self._get_neighbors(node_idx, adj)
            nodes_to_enhance.update(neighbors)
        
        nodes_to_enhance = list(nodes_to_enhance)
        
        for node_idx in nodes_to_enhance:
            try:
                related_vectors, similarities = self.get_related_feature_vectors(
                    node_id=node_idx,
                    raw_texts=raw_texts,
                    k=k,
                    filter_datasets=filter_datasets
                )
                
                if not related_vectors:
                    self.logger.warning(f"node {node_idx} not found related feature vectors, skip")
                    continue
                
                weighted_vector = self._compute_weighted_features(related_vectors, similarities)
                
                if weighted_vector is not None:
                    if weighted_vector.device != enhanced_features.device:
                        weighted_vector = weighted_vector.to(enhanced_features.device)
                    
                    enhanced_features[node_idx] += rag_weight * weighted_vector
                    
                
            except Exception as e:
                self.logger.error(f"node {node_idx} enhancement failed: {e}")
                continue
        
        return enhanced_features
    
    
    def _compute_weighted_features(self, 
                                 related_vectors: List[torch.Tensor], 
                                 similarities: List[float]) -> Optional[torch.Tensor]:

        if not related_vectors or not similarities:
            return None
        
        if len(related_vectors) != len(similarities):
            self.logger.error(f"feature vector number({len(related_vectors)}) does not match similarity number({len(similarities)})")
            return None
        
        try:
            weights = torch.tensor(similarities, device=related_vectors[0].device)
            weights = torch.softmax(weights, dim=0)
            
            weighted_vector = torch.zeros_like(related_vectors[0])
            for i, vector in enumerate(related_vectors):
                weighted_vector += weights[i] * vector
            
            return weighted_vector
            
        except Exception as e:
            self.logger.error(f"failed to compute weighted features: {e}")
            return None
    
    def get_pretrain_feature_vector(self, dataset_name: str, node_id: int) -> Optional[torch.Tensor]:

        if not hasattr(self, 'pretrain_features') or self.pretrain_features is None:
            self.logger.warning("pretrained feature matrix not loaded")
            return None
        
        if dataset_name not in self.pretrain_features:
            self.logger.warning(f"dataset {dataset_name} not in pretrained features")
            return None
        
        features = self.pretrain_features[dataset_name]
        
        if node_id >= features.shape[0] or node_id < 0:
            self.logger.warning(f"node ID {node_id} out of range [0, {features.shape[0]-1}]")
            return None
        
        return features[node_id]
    
    def get_motif_enhanced_feature(self, 
                                 node_idx: int,
                                 edge_index: torch.Tensor,
                                 num_nodes: int,
                                 search_domain: str,
                                 dataset_name: str = None) -> Optional[torch.Tensor]:

        if dataset_name is None:
            dataset_name = self.config.dataset_name
            
        try:
            if not self._are_components_initialized():
                self.logger.error("Motif retriever not initialized, please call initialize_components()")
                return None
            
            query_subgraph = self._extract_1hop_subgraph(
                center_node_idx=node_idx,
                edge_index=edge_index,
                num_nodes=num_nodes,
                dataset_name=dataset_name
            )
            
            motif_results = self.motif_retriever.search(
                query_subgraph=query_subgraph,
                search_domain=search_domain,
                k=1 
            )
            
            if not motif_results:
                self.logger.warning(f"node {node_idx} in domain {search_domain} not found similar motif")
                return None
            
            best_result = motif_results[0]
            metadata = best_result.get('metadata', {})
            
            domain_name = metadata.get('domain', '')
            center_node_original_idx = metadata.get('center_node_original_idx')
            
            if not domain_name or center_node_original_idx is None:
                self.logger.warning(f"Motif result metadata incomplete: {metadata}")
                return None
            
            pretrain_feature = self.get_pretrain_feature_vector(
                dataset_name=domain_name,
                node_id=center_node_original_idx
            )
            
            if pretrain_feature is None:
                self.logger.warning(f"node {domain_name}:{center_node_original_idx} not found pretrained feature")
                return None
            
            return pretrain_feature
            
        except Exception as e:
            self.logger.error(f"Motif retrieval failed: {e}")
            return None
    
    def get_multi_domain_motif_enhanced_feature(self,
                                              node_idx: int,
                                              edge_index: torch.Tensor,
                                              num_nodes: int,
                                              pretrain_dataset_names: Set[str],
                                              pretrain_weights: Dict[str, torch.Tensor],
                                              dataset_name: str = None,
                                              motif_weight: float = 0.15) -> Optional[torch.Tensor]:

        if dataset_name is None:
            dataset_name = self.config.dataset_name
            
        try:
            if not self._are_components_initialized():
                self.logger.error("Motif retriever not initialized, please call initialize_components()")
                return None
            
            domain_features = []
            domain_weights = []
            successful_domains = []
            
            for domain in pretrain_dataset_names:
                try:
                    motif_feature = self.get_motif_enhanced_feature(
                        node_idx=node_idx,
                        edge_index=edge_index,
                        num_nodes=num_nodes,
                        search_domain=domain,
                        dataset_name=dataset_name
                    )
                    
                    if motif_feature is not None:
                        if domain in pretrain_weights:
                            weight = pretrain_weights[domain]
                            if isinstance(weight, torch.Tensor):
                                if weight.shape == (1, self.config.feature_dim):
                                    weight = weight.squeeze(0)  
                                elif weight.shape == (self.config.feature_dim,):
                                    pass  
                                else:
                                    if weight.numel() == self.config.feature_dim:
                                        weight = weight.view(self.config.feature_dim)
                                    else:
                                        weight = torch.full((self.config.feature_dim,), weight.mean().item(), device=weight.device)
                            else:
                                weight = torch.ones(self.config.feature_dim, device=motif_feature.device)
                        else:
                            weight = torch.ones(self.config.feature_dim, device=motif_feature.device)
                            self.logger.warning(f"dataset {domain} not found pretrained weights, using all 1 weights")
                        
                        domain_features.append(motif_feature)
                        domain_weights.append(weight)
                        successful_domains.append(domain)
                        
                    else:
                        pass
                        
                except Exception as e:
                    self.logger.error(f"failed to get {domain} domain feature: {e}")
                    continue
            
            if not domain_features:
                self.logger.error("all pretrained datasets failed to get Motif enhanced feature")
                return None
            
            feature_dims = [feat.shape[0] for feat in domain_features]
            if len(set(feature_dims)) > 1:
                self.logger.warning(f"feature dimension inconsistent, will truncate to minimum dimension")
                min_dim = min(feature_dims)
                domain_features = [feat[:min_dim] for feat in domain_features]
                domain_weights = [w[:min_dim] for w in domain_weights]
            
            features_tensor = torch.stack(domain_features) 
            weights_tensor = torch.stack(domain_weights) 
            
            weights_normalized = torch.softmax(weights_tensor, dim=0) 
            
            weighted_feature = torch.sum(weights_normalized * features_tensor, dim=0) 
            
            
            return weighted_feature
            
        except Exception as e:
            self.logger.error(f"multi-domain Motif feature retrieval failed: {e}")
            return None