salim0986/graph-bug-ai

1"""2Context Merger3Combines temporary (PR-only) and permanent (main branch) GraphRAG contexts4"""5 6from typing import Dict, List, Any, Optional7from dataclasses import dataclass8from .logger import setup_logger9from .temporary_graph import TemporaryGraphBuilder, TemporaryVectorBuilder, TempNode10 11logger = setup_logger(__name__)12 13 14@dataclass15class MergedContext:16    """Unified context combining temporary and permanent data"""17    18    # Node-level context19    dependencies: List[Dict[str, Any]]  # Functions/classes this depends on20    dependents: List[Dict[str, Any]]    # Functions/classes that depend on this21    similar_code: List[Dict[str, Any]]  # Similar code snippets22    23    # File-level context24    file_imports: List[str]25    file_dependencies: List[str]26    27    # Metadata28    source: str  # "temporary", "permanent", or "merged"29    temp_nodes_count: int = 030    permanent_nodes_count: int = 031 32 33class ContextMerger:34    """35    Merges temporary and permanent GraphRAG contexts36    37    Strategy:38    1. Check permanent database first (Neo4j + Qdrant)39    2. Check temporary in-memory graph second40    3. Combine results with deduplication41    4. Prioritize temporary results (newer code)42    """43    44    def __init__(45        self,46        temp_graph: Optional[TemporaryGraphBuilder] = None,47        temp_vector: Optional[TemporaryVectorBuilder] = None48    ):49        self.temp_graph = temp_graph50        self.temp_vector = temp_vector51        52    def merge_file_context(53        self,54        filename: str,55        permanent_context: Dict[str, Any],56        code_content: Optional[str] = None57    ) -> MergedContext:58        """59        Merge context for a specific file60        61        Args:62            filename: File path63            permanent_context: Context from permanent databases (Neo4j + Qdrant)64            code_content: Optional raw file content for on-the-fly parsing65            66        Returns:67            MergedContext with combined data68        """69        logger.info(f"[ContextMerger] Merging context for {filename}")70        71        # Extract permanent data72        perm_deps = permanent_context.get("dependencies", [])73        perm_dependents = permanent_context.get("dependents", [])74        perm_similar = permanent_context.get("similar_code", [])75        perm_imports = permanent_context.get("imports", [])76        perm_file_deps = permanent_context.get("file_dependencies", [])77        78        # Initialize temporary data containers79        temp_deps = []80        temp_dependents = []81        temp_similar = []82        temp_imports = []83        temp_file_deps = []84        85        # Check if file exists in temporary graph86        has_temp_data = False87        if self.temp_graph and filename in self.temp_graph.files:88            has_temp_data = True89            file_node = self.temp_graph.files[filename]90            91            # Get temporary imports and dependencies92            temp_imports = list(file_node.imports)93            temp_file_deps = self.temp_graph.get_file_dependencies(filename)94            95            # Get node-level dependencies from temporary graph96            for node in file_node.nodes:97                # Dependencies (what this node calls)98                node_deps = self.temp_graph.get_node_dependencies(node.id)99                temp_deps.extend(node_deps)100                101                # Similar code from temporary vectors102                if self.temp_vector:103                    similar = self.temp_vector.find_similar_to_node(104                        node.id,105                        limit=5,106                        min_score=0.7107                    )108                    temp_similar.extend(similar)109        110        # Merge with deduplication111        merged_dependencies = self._merge_lists(perm_deps, temp_deps, key="name")112        merged_dependents = self._merge_lists(perm_dependents, temp_dependents, key="name")113        merged_similar = self._merge_lists(perm_similar, temp_similar, key="node_id")114        merged_imports = list(set(perm_imports + temp_imports))115        merged_file_deps = list(set(perm_file_deps + temp_file_deps))116        117        # Determine source118        if has_temp_data and perm_deps:119            source = "merged"120        elif has_temp_data:121            source = "temporary"122        else:123            source = "permanent"124        125        return MergedContext(126            dependencies=merged_dependencies,127            dependents=merged_dependents,128            similar_code=merged_similar,129            file_imports=merged_imports,130            file_dependencies=merged_file_deps,131            source=source,132            temp_nodes_count=len(temp_deps) + len(temp_similar),133            permanent_nodes_count=len(perm_deps) + len(perm_similar)134        )135    136    def merge_similar_code_search(137        self,138        query: str,139        permanent_results: List[Dict[str, Any]],140        limit: int = 10141    ) -> List[Dict[str, Any]]:142        """143        Merge similar code search results from permanent and temporary144        145        Args:146            query: Search query147            permanent_results: Results from Qdrant permanent database148            limit: Maximum results to return149            150        Returns:151            Merged list of similar code snippets152        """153        temp_results = []154        155        # Search temporary vectors if available156        if self.temp_vector:157            temp_results = self.temp_vector.search_similar(158                query,159                limit=limit,160                min_score=0.7161            )162            # Mark as temporary source163            for result in temp_results:164                result["source"] = "temporary"165        166        # Mark permanent results167        for result in permanent_results:168            result["source"] = "permanent"169        170        # Combine and sort by score171        all_results = permanent_results + temp_results172        all_results.sort(key=lambda x: x.get("score", 0), reverse=True)173        174        # Deduplicate by text similarity175        deduplicated = self._deduplicate_similar_code(all_results)176        177        return deduplicated[:limit]178    179    def _merge_lists(180        self,181        list1: List[Dict[str, Any]],182        list2: List[Dict[str, Any]],183        key: str184    ) -> List[Dict[str, Any]]:185        """186        Merge two lists of dicts with deduplication by key187        188        Strategy: Prioritize items from list2 (temporary) over list1 (permanent)189        """190        seen = {}191        result = []192        193        # Add list2 first (temporary has priority)194        for item in list2:195            item_key = item.get(key)196            if item_key and item_key not in seen:197                seen[item_key] = True198                item["source"] = "temporary"199                result.append(item)200        201        # Add list1 (permanent) for items not in list2202        for item in list1:203            item_key = item.get(key)204            if item_key and item_key not in seen:205                seen[item_key] = True206                item["source"] = "permanent"207                result.append(item)208        209        return result210    211    def _deduplicate_similar_code(212        self,213        results: List[Dict[str, Any]]214    ) -> List[Dict[str, Any]]:215        """Remove duplicate similar code entries based on text similarity"""216        if not results:217            return []218        219        deduplicated = [results[0]]220        221        for result in results[1:]:222            # Check if this result is too similar to any existing result223            is_duplicate = False224            for existing in deduplicated:225                if self._are_similar_texts(226                    result.get("text", ""),227                    existing.get("text", "")228                ):229                    is_duplicate = True230                    break231            232            if not is_duplicate:233                deduplicated.append(result)234        235        return deduplicated236    237    def _are_similar_texts(self, text1: str, text2: str, threshold: float = 0.8) -> bool:238        """Check if two texts are very similar (simple Jaccard similarity)"""239        if not text1 or not text2:240            return False241        242        # Simple word-based similarity243        words1 = set(text1.lower().split())244        words2 = set(text2.lower().split())245        246        if not words1 or not words2:247            return False248        249        intersection = len(words1 & words2)250        union = len(words1 | words2)251        252        return (intersection / union) >= threshold if union > 0 else False253    254    def get_statistics(self) -> Dict[str, Any]:255        """Get statistics about temporary and permanent contexts"""256        stats = {257            "has_temporary_graph": self.temp_graph is not None,258            "has_temporary_vectors": self.temp_vector is not None259        }260        261        if self.temp_graph:262            stats.update({263                "temp_files": len(self.temp_graph.files),264                "temp_nodes": len(self.temp_graph.nodes)265            })266        267        if self.temp_vector:268            stats.update({269                "temp_vectors": len(self.temp_vector.vectors)270            })271        272        return stats273