GraphRAG Tutorial

# Install required packages
pip install faiss-cpu langchain langchain-openai matplotlib networkx nltk numpy python-dotenv scikit-learn spacy tqdm

import networkx as nx
import heapq
import numpy as np
import spacy
import nltk
import os
from typing import List, Tuple, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

from langchain.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.callbacks import get_openai_callback
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.document_loaders import PyPDFLoader

from sklearn.metrics.pairwise import cosine_similarity
from pydantic import BaseModel, Field
from nltk.stem import WordNetLemmatizer
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')

# Download NLTK data
nltk.download('punkt', quiet=True)
nltk.download('wordnet', quiet=True)

class DocumentProcessor:
    def __init__(self):
        """Initialize with text splitter and embeddings."""
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000, 
            chunk_overlap=200
        )
        self.embeddings = OpenAIEmbeddings()

    def process_documents(self, documents):
        """Split documents and create vector store."""
        splits = self.text_splitter.split_documents(documents)
        vector_store = FAISS.from_documents(splits, self.embeddings)
        return splits, vector_store

    def create_embeddings_batch(self, texts, batch_size=32):
        """Create embeddings in batches for efficiency."""
        embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i+batch_size]
            batch_embeddings = self.embeddings.embed_documents(batch)
            embeddings.extend(batch_embeddings)
        return np.array(embeddings)

class KnowledgeGraph:
    def __init__(self):
        self.graph = nx.Graph()
        self.lemmatizer = WordNetLemmatizer()
        self.concept_cache = {}
        self.nlp = self._load_spacy_model()
        self.edges_threshold = 0.8  # Similarity threshold for edges

    def _load_spacy_model(self):
        """Load spaCy model for NER."""
        try:
            return spacy.load("en_core_web_sm")
        except OSError:
            from spacy.cli import download
            download("en_core_web_sm")
            return spacy.load("en_core_web_sm")

    def build_graph(self, splits, llm, embedding_model):
        """Build the complete knowledge graph."""
        self._add_nodes(splits)
        embeddings = self._create_embeddings(splits, embedding_model)
        self._extract_concepts(splits, llm)
        self._add_edges(embeddings)

    def _add_nodes(self, splits):
        """Create nodes from document chunks."""
        for i, split in enumerate(splits):
            self.graph.add_node(i, content=split.page_content)

    def _create_embeddings(self, splits, embedding_model):
        texts = [split.page_content for split in splits]
        return embedding_model.embed_documents(texts)

    def _extract_concepts_and_entities(self, content, llm):
        """Extract concepts using spaCy NER + LLM."""
        if content in self.concept_cache:
            return self.concept_cache[content]
        
        # Extract named entities with spaCy
        doc = self.nlp(content)
        named_entities = [
            ent.text for ent in doc.ents 
            if ent.label_ in ["PERSON", "ORG", "GPE", "WORK_OF_ART"]
        ]
        
        # Extract general concepts with LLM
        concept_prompt = PromptTemplate(
            input_variables=["text"],
            template="Extract key concepts from:\n\n{text}\n\nKey concepts:"
        )
        concept_chain = concept_prompt | llm.with_structured_output(Concepts)
        general_concepts = concept_chain.invoke({"text": content}).concepts_list
        
        all_concepts = list(set(named_entities + general_concepts))
        self.concept_cache[content] = all_concepts
        return all_concepts

    def _extract_concepts(self, splits, llm):
        """Extract concepts with multi-threading."""
        with ThreadPoolExecutor() as executor:
            future_to_node = {
                executor.submit(self._extract_concepts_and_entities, split.page_content, llm): i 
                for i, split in enumerate(splits)
            }
            for future in tqdm(as_completed(future_to_node), total=len(splits)):
                node = future_to_node[future]
                self.graph.nodes[node]['concepts'] = future.result()

    def _add_edges(self, embeddings):
        """Add edges based on similarity and shared concepts."""
        similarity_matrix = cosine_similarity(embeddings)
        num_nodes = len(self.graph.nodes)
        
        for node1 in tqdm(range(num_nodes), desc="Adding edges"):
            for node2 in range(node1 + 1, num_nodes):
                similarity = similarity_matrix[node1][node2]
                if similarity > self.edges_threshold:
                    shared = set(self.graph.nodes[node1]['concepts']) & \
                             set(self.graph.nodes[node2]['concepts'])
                    weight = self._calculate_edge_weight(node1, node2, similarity, shared)
                    self.graph.add_edge(node1, node2, weight=weight, similarity=similarity)

    def _calculate_edge_weight(self, node1, node2, similarity, shared, alpha=0.7, beta=0.3):
        """Weight = α×similarity + β×normalized_shared_concepts"""
        max_shared = min(
            len(self.graph.nodes[node1]['concepts']), 
            len(self.graph.nodes[node2]['concepts'])
        )
        normalized = len(shared) / max_shared if max_shared > 0 else 0
        return alpha * similarity + beta * normalized

    def _lemmatize_concept(self, concept):
        return ' '.join([self.lemmatizer.lemmatize(w) for w in concept.lower().split()])

class QueryEngine:
    def __init__(self, vector_store, knowledge_graph, llm):
        self.vector_store = vector_store
        self.knowledge_graph = knowledge_graph
        self.llm = llm
        self.answer_check_chain = self._create_answer_check_chain()

    def _create_answer_check_chain(self):
        prompt = PromptTemplate(
            input_variables=["query", "context"],
            template="""Given the query: '{query}'
And the context:
{context}

Does this provide a complete answer? If yes, provide it."""
        )
        return prompt | self.llm.with_structured_output(AnswerCheck)

    def query(self, query: str):
        """Process query with graph traversal."""
        with get_openai_callback() as cb:
            relevant_docs = self._retrieve_relevant_documents(query)
            context, path, content, answer = self._expand_context(query, relevant_docs)
            
            if not answer:
                response_prompt = PromptTemplate(
                    input_variables=["query", "context"],
                    template="Context: {context}\n\nQuery: {query}\n\nAnswer:"
                )
                answer = (response_prompt | self.llm).invoke({
                    "query": query, "context": context
                })
            
            print(f"Total Tokens: {cb.total_tokens}")
            print(f"Cost: ${cb.total_cost}")
        
        return answer, path, content

    def _retrieve_relevant_documents(self, query: str):
        """Retrieve with contextual compression."""
        retriever = self.vector_store.as_retriever(search_kwargs={"k": 5})
        compressor = LLMChainExtractor.from_llm(self.llm)
        compression_retriever = ContextualCompressionRetriever(
            base_compressor=compressor, base_retriever=retriever
        )
        return compression_retriever.invoke(query)

    def _expand_context(self, query: str, relevant_docs):
        """Dijkstra-like graph traversal for context expansion."""
        expanded_context = ""
        traversal_path = []
        visited_concepts = set()
        filtered_content = {}
        final_answer = ""
        
        priority_queue = []
        distances = {}
        
        # Initialize with relevant document nodes
        for doc in relevant_docs:
            closest = self.vector_store.similarity_search_with_score(doc.page_content, k=1)
            node_content, score = closest[0]
            node = next(
                n for n in self.knowledge_graph.graph.nodes 
                if self.knowledge_graph.graph.nodes[n]['content'] == node_content.page_content
            )
            priority = 1 / score
            heapq.heappush(priority_queue, (priority, node))
            distances[node] = priority
        
        # Traverse graph
        while priority_queue:
            current_priority, current_node = heapq.heappop(priority_queue)
            
            if current_priority > distances.get(current_node, float('inf')):
                continue
            
            if current_node not in traversal_path:
                traversal_path.append(current_node)
                node_content = self.knowledge_graph.graph.nodes[current_node]['content']
                node_concepts = self.knowledge_graph.graph.nodes[current_node]['concepts']
                
                filtered_content[current_node] = node_content
                expanded_context += "\n" + node_content if expanded_context else node_content
                
                # Check if we have complete answer
                is_complete, answer = self._check_answer(query, expanded_context)
                if is_complete:
                    final_answer = answer
                    break
                
                # Process concepts and explore neighbors
                node_concepts_set = set(
                    self.knowledge_graph._lemmatize_concept(c) for c in node_concepts
                )
                if not node_concepts_set.issubset(visited_concepts):
                    visited_concepts.update(node_concepts_set)
                    
                    for neighbor in self.knowledge_graph.graph.neighbors(current_node):
                        edge_weight = self.knowledge_graph.graph[current_node][neighbor]['weight']
                        distance = current_priority + (1 / edge_weight)
                        
                        if distance < distances.get(neighbor, float('inf')):
                            distances[neighbor] = distance
                            heapq.heappush(priority_queue, (distance, neighbor))
        
        return expanded_context, traversal_path, filtered_content, final_answer

    def _check_answer(self, query, context):
        response = self.answer_check_chain.invoke({"query": query, "context": context})
        return response.is_complete, response.answer

import matplotlib.pyplot as plt
import matplotlib.patches as patches

class Visualizer:
    @staticmethod
    def visualize_traversal(graph, traversal_path):
        """Visualize the graph traversal path."""
        fig, ax = plt.subplots(figsize=(16, 12))
        pos = nx.spring_layout(graph, k=1, iterations=50)
        
        # Draw edges with weight colors
        edges = graph.edges()
        weights = [graph[u][v].get('weight', 0.5) for u, v in edges]
        nx.draw_networkx_edges(graph, pos, edge_color=weights, 
                               edge_cmap=plt.cm.Blues, width=2, ax=ax)
        
        # Draw nodes
        nx.draw_networkx_nodes(graph, pos, node_color='lightblue', 
                               node_size=3000, ax=ax)
        
        # Draw traversal path
        for i in range(len(traversal_path) - 1):
            start, end = traversal_path[i], traversal_path[i + 1]
            arrow = patches.FancyArrowPatch(
                pos[start], pos[end],
                connectionstyle="arc3,rad=0.3",
                color='red', arrowstyle="->",
                mutation_scale=20, linestyle='--', linewidth=2
            )
            ax.add_patch(arrow)
        
        # Highlight start/end nodes
        nx.draw_networkx_nodes(graph, pos, nodelist=[traversal_path[0]], 
                               node_color='lightgreen', node_size=3000, ax=ax)
        nx.draw_networkx_nodes(graph, pos, nodelist=[traversal_path[-1]], 
                               node_color='lightcoral', node_size=3000, ax=ax)
        
        # Labels
        labels = {n: graph.nodes[n].get('concepts', [''])[0][:20] for n in graph.nodes()}
        nx.draw_networkx_labels(graph, pos, labels, font_size=8, ax=ax)
        
        ax.set_title("GraphRAG Traversal Visualization")
        ax.axis('off')
        plt.tight_layout()
        plt.show()

class GraphRAG:
    def __init__(self):
        self.llm = ChatOpenAI(temperature=0, model_name="gpt-4o-mini", max_tokens=4000)
        self.embedding_model = OpenAIEmbeddings()
        self.document_processor = DocumentProcessor()
        self.knowledge_graph = KnowledgeGraph()
        self.query_engine = None
        self.visualizer = Visualizer()

    def process_documents(self, documents):
        """Process documents and build knowledge graph."""
        splits, vector_store = self.document_processor.process_documents(documents)
        self.knowledge_graph.build_graph(splits, self.llm, self.embedding_model)
        self.query_engine = QueryEngine(vector_store, self.knowledge_graph, self.llm)

    def query(self, query: str):
        """Query with visualization."""
        response, traversal_path, _ = self.query_engine.query(query)
        
        if traversal_path:
            self.visualizer.visualize_traversal(self.knowledge_graph.graph, traversal_path)
        
        return response

# Load documents
loader = PyPDFLoader("data/Understanding_Climate_Change.pdf")
documents = loader.load()[:10]  # Use first 10 pages for demo

# Initialize GraphRAG
graph_rag = GraphRAG()

# Process documents (builds knowledge graph)
graph_rag.process_documents(documents)

# Query the system
query = "What is the main cause of climate change?"
response = graph_rag.query(query)

print(f"Answer: {response}")