Merge pull request #57 from data-exp-lab/fe-rag

GraphRAG Implementation

Author: yilinxia

.gitignore

@@ -11,6 +11,60 @@ dist
 *temp
 *repo_metadata.json
 __pycache__
+
+# Database files
 *.duckdb
 *.gexf
-*.gexf.gz
+*.gexf.gz
+
+# GraphRAG specific
+backend/app/kuzu/
+backend/app/cache/
+rag/kuzu_db_*
+*.kuzu
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+venv/
+env/
+ENV/
+env.bak/
+venv.bak/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db

backend/app/main.py

@@ -1,13 +1,15 @@
-from flask import Flask, jsonify, request, send_file, url_for
+from flask import Flask, jsonify, request, send_file, url_for, Response, stream_template
 from flask_cors import CORS
 from services.topic_service import TopicService
 from services.ai_service import AITopicProcessor
 from services.gexf_node_service import GexfNodeGenerator
 from services.edge_generation_service import EdgeGenerationService
+from services.graphrag_service import graphrag_service
 import os
 import asyncio
 import re
 import json
+import time
 
 app = Flask(__name__, static_folder='gexf', static_url_path='/gexf')
 CORS(
@@ -26,6 +28,43 @@
 gexf_node_service = GexfNodeGenerator()
 edge_generation_service = EdgeGenerationService()
 
+# Global progress tracking for GraphRAG setup
+graphrag_progress = {
+    "current_step": "",
+    "current": 0,
+    "total": 0,
+    "message": "",
+    "status": "idle"  # idle, running, completed, error
+}
+
+# Global variable to track if GraphRAG is set up
+graphrag_ready = False
+
+@app.route("/api/graphrag-health", methods=["GET"])
+def graphrag_health():
+    """Check if GraphRAG backend is ready and set up."""
+    global graphrag_ready
+    try:
+        if graphrag_ready:
+            return jsonify({
+                "success": True,
+                "ready": True,
+                "message": "GraphRAG backend is ready"
+            })
+        else:
+            return jsonify({
+                "success": True,
+                "ready": False,
+                "message": "GraphRAG backend is not set up"
+            }), 503
+    except Exception as e:
+        return jsonify({
+            "success": False,
+            "ready": False,
+            "error": str(e),
+            "message": "Error checking GraphRAG health"
+        }), 500
+
 
 @app.route("/api/process-topics", methods=["GET", "POST"])
 def process_topics():
@@ -524,15 +563,16 @@ def create_edges_on_graph():
             "shared_organization_enabled": false,
             "common_stargazers_enabled": true,
             "stargazer_threshold": 5,
-            "use_and_logic": false,
-            "strict_and_logic": true
-        }
+            "use_and_logic": false
+        },
+        "filtered_node_ids": ["node1", "node2", ...] // Optional: only consider these nodes for edge creation
     }
     """
     try:
         data = request.get_json()
         gexf_content = data.get("gexfContent", "")
         criteria_config = data.get("criteria_config", {})
+        filtered_node_ids = data.get("filtered_node_ids", None)  # Get filtered node IDs
 
         if not gexf_content:
             return jsonify({
@@ -582,7 +622,7 @@ def create_edges_on_graph():
         edge_service = EdgeGenerationService()
 
         # Create edges based on the criteria
-        edges_created = edge_service.create_edges_on_existing_graph(G, criteria_config)
+        edges_created = edge_service.create_edges_on_existing_graph(G, criteria_config, filtered_node_ids)
 
         # Save the updated graph
         import hashlib
@@ -627,6 +667,153 @@ def create_edges_on_graph():
         }), 500
 
 
+@app.route("/api/graphrag-reset-progress", methods=["POST", "OPTIONS"])
+def graphrag_reset_progress_endpoint():
+    """Reset GraphRAG progress status to initial state."""
+    if request.method == "OPTIONS":
+        return "", 200
+    
+    global graphrag_progress
+    graphrag_progress = {
+        "current_step": "Initializing...",
+        "current": 0,
+        "total": 100,
+        "message": "Preparing GraphRAG setup",
+        "status": "running"
+    }
+    return jsonify({"success": True, "message": "Progress reset"})
+
+@app.route("/api/graphrag-progress", methods=["GET"])
+def graphrag_progress_endpoint():
+    """Server-Sent Events endpoint for GraphRAG progress updates."""
+    def generate():
+        while True:
+            # Send current progress
+            data = f"data: {json.dumps(graphrag_progress)}\n\n"
+            yield data
+            
+            # If completed or error, stop streaming
+            if graphrag_progress["status"] in ["completed", "error"]:
+                break
+            
+            time.sleep(0.5)  # Update every 0.5 seconds for more responsive updates
+    
+    return Response(generate(), mimetype="text/event-stream")
+
+
+@app.route("/api/graphrag-setup", methods=["POST"])
+def graphrag_setup_endpoint():
+    """GraphRAG setup endpoint with progress tracking."""
+    global graphrag_progress, graphrag_ready
+    
+    try:
+        data = request.get_json()
+        
+        # Extract parameters
+        provider = data.get("provider", "openai")
+        api_keys = data.get("apiKeys", {})
+        graph_file = data.get("graphFile", "")
+        
+        if not graph_file:
+            return jsonify({
+                "success": False,
+                "error": "Graph file is required",
+                "message": "Please provide a graph file"
+            }), 400
+        
+        github_token = api_keys.get("githubToken", "")
+        if not github_token:
+            return jsonify({
+                "success": False,
+                "error": "GitHub token is required",
+                "message": "Please provide a GitHub personal access token"
+            }), 400
+        
+        # Reset progress
+        graphrag_progress = {
+            "current_step": "Starting setup...",
+            "current": 0,
+            "total": 100,
+            "message": "Initializing GraphRAG setup",
+            "status": "running"
+        }
+        
+        # Setup database from GEXF content with progress updates
+        setup_result = graphrag_service.setup_database_from_gexf_with_progress(graph_file, github_token, graphrag_progress)
+        if not setup_result["success"]:
+            graphrag_progress["status"] = "error"
+            graphrag_progress["message"] = setup_result.get("error", "Setup failed")
+            return jsonify(setup_result), 500
+        
+        # Initialize GraphRAG with the selected provider
+        graphrag_progress["current_step"] = "Initializing AI system..."
+        graphrag_progress["current"] = 90
+        graphrag_progress["message"] = "Setting up AI analysis system"
+        
+        init_result = graphrag_service.initialize_graphrag(provider, api_keys)
+        if not init_result["success"]:
+            graphrag_progress["status"] = "error"
+            graphrag_progress["message"] = init_result.get("error", "AI initialization failed")
+            return jsonify(init_result), 500
+        
+        # Mark as completed and set ready flag
+        graphrag_progress["status"] = "completed"
+        graphrag_progress["current"] = 100
+        graphrag_progress["message"] = "GraphRAG setup completed successfully!"
+        graphrag_ready = True
+        
+        return jsonify({
+            "success": True,
+            "message": "GraphRAG setup completed successfully",
+            "ready": True
+        })
+        
+    except Exception as e:
+        graphrag_progress["status"] = "error"
+        graphrag_progress["message"] = str(e)
+        return jsonify({
+            "success": False,
+            "error": str(e),
+            "message": "An error occurred during GraphRAG setup"
+        }), 500
+
+
+@app.route("/api/graphrag", methods=["POST"])
+def graphrag_endpoint():
+    """GraphRAG endpoint for AI-powered graph analysis."""
+    try:
+        data = request.get_json()
+        
+        # Extract parameters
+        query = data.get("query", "")
+        provider = data.get("provider", "openai")
+        api_keys = data.get("apiKeys", {})
+        
+        if not query:
+            return jsonify({
+                "success": False,
+                "error": "Query is required",
+                "message": "Please provide a query"
+            }), 400
+        
+        # Execute the query
+        query_result = graphrag_service.query_graphrag(query)
+        if not query_result["success"]:
+            return jsonify(query_result), 500
+        
+        return jsonify({
+            "success": True,
+            "result": query_result["result"]
+        })
+        
+    except Exception as e:
+        return jsonify({
+            "success": False,
+            "error": str(e),
+            "message": "An error occurred while processing the GraphRAG query"
+        }), 500
+
+
 @app.route("/")
 def home():
     return "Hello World!"

backend/app/services/edge_generation_service.py

@@ -148,13 +148,14 @@ def generate_edges_with_criteria(
 
         return G, edge_stats
 
-    def create_edges_on_existing_graph(self, G: nx.Graph, criteria_config: Dict[str, any]) -> Dict[str, any]:
+    def create_edges_on_existing_graph(self, G: nx.Graph, criteria_config: Dict[str, any], filtered_node_ids: List[str] = None) -> Dict[str, any]:
         """
         Create edges on an existing graph based on specified criteria.
         
         Args:
             G: Existing NetworkX graph with nodes
             criteria_config: Configuration for edge generation criteria
+            filtered_node_ids: Optional list of node IDs to consider for edge creation (only these nodes will be used)
             
         Returns:
             Dictionary with statistics about created edges
@@ -184,10 +185,27 @@ def create_edges_on_existing_graph(self, G: nx.Graph, criteria_config: Dict[str,
         # Remove all existing edges first
         G.remove_edges_from(list(G.edges()))
 
-        # Get all nodes
-        nodes = list(G.nodes())
-        if len(nodes) < 2:
-            return {'message': 'Not enough nodes to create edges'}
+        # Use filtered nodes if provided, otherwise use all nodes
+        if filtered_node_ids is not None:
+            # Validate that all filtered node IDs exist in the graph
+            existing_nodes = set(G.nodes())
+            valid_filtered_nodes = [node_id for node_id in filtered_node_ids if node_id in existing_nodes]
+            
+            if len(valid_filtered_nodes) < 2:
+                return {
+                    'message': f'Not enough valid filtered nodes to create edges. Found {len(valid_filtered_nodes)} valid nodes out of {len(filtered_node_ids)} provided.',
+                    'total_edges': 0,
+                    'total_nodes': len(valid_filtered_nodes)
+                }
+            
+            nodes = valid_filtered_nodes
+            print(f"Using {len(nodes)} filtered nodes for edge creation out of {len(existing_nodes)} total nodes")
+        else:
+            # Use all nodes if no filtering is applied
+            nodes = list(G.nodes())
+            if len(nodes) < 2:
+                return {'message': 'Not enough nodes to create edges'}
+            print(f"Using all {len(nodes)} nodes for edge creation")
 
         # Generate edges based on enabled criteria
         edge_stats = {}