Graph RAG Techniques

calendarDec 12, 2024 · 8 min read · ai-knowhow graph-rag knowledge-graph llm neo4j  ·
Share on: linkedincopy

Graph-based Retrieval-Augmented Generation for enhanced context and relationship understanding.

Core Idea

Graph RAG represents knowledge as a structured graph (entities and relationships) rather than flat document chunks, enabling multi-hop reasoning and relationship-aware retrieval.

Mathematical Foundation

Traditional vector RAG uses point-wise similarity: $$\text{retrieve}(q) = \arg\max_{d \in D} \text{sim}(\mathbf{q}, \mathbf{d})$$

Graph RAG uses graph traversal to retrieve connected subgraphs: $$\text{retrieve}(q) = \text{traverse}(G, E_q, h)$$

where:

  • $G = (V, E)$ is the knowledge graph with vertices (entities) and edges (relationships)
  • $E_q$ are entities extracted from query $q$
  • $h$ is the hop depth for traversal
  • The traversal returns a subgraph $G' \subseteq G$ containing relevant entities and their relationships

Key Advantages:

  1. Relationship Preservation: Captures explicit connections: $(e_1, r, e_2) \in E$
  2. Multi-hop Reasoning: Traverses paths: $e_1 \xrightarrow{r_1} e_2 \xrightarrow{r_2} e_3$ to answer complex queries
  3. Structure Awareness: Maintains document hierarchy and entity relationships
  4. Hybrid Retrieval: Combines graph traversal with vector similarity for entities: $\text{hybrid}(q) = \alpha \cdot \text{graph}(q) + (1-\alpha) \cdot \text{vector}(q)$

Process:

  1. Extraction: Extract entities and relationships from documents: ${(e_i, r_{ij}, e_j)}$
  2. Graph Construction: Build knowledge graph $G$ with nodes and edges
  3. Query Processing: Identify query entities $E_q$ and traverse connected subgraph
  4. Context Formation: Convert subgraph to textual context preserving relationships
  5. Generation: LLM generates response using graph-structured context

This enables answering questions like "Who collaborated with authors of papers about RAG?" which requires following relationships: $\text{Paper} \xrightarrow{\text{hasAuthor}} \text{Author} \xrightarrow{\text{collaboratesWith}} \text{Author}$.

Graph RAG Architecture

Why Graph RAG?

Advantages over Vector RAG:

  • Captures relationships between entities
  • Enables multi-hop reasoning
  • Preserves document structure
  • Better for complex queries requiring context

Building Knowledge Graph

Entity and Relationship Extraction

 1from langchain.llms import OpenAI
 2from langchain.chains import create_extraction_chain
 3
 4schema = {
 5    "properties": {
 6        "entities": {
 7            "type": "array",
 8            "items": {
 9                "type": "object",
10                "properties": {
11                    "name": {"type": "string"},
12                    "type": {"type": "string"},
13                }
14            }
15        },
16        "relationships": {
17            "type": "array",
18            "items": {
19                "type": "object",
20                "properties": {
21                    "source": {"type": "string"},
22                    "target": {"type": "string"},
23                    "type": {"type": "string"},
24                }
25            }
26        }
27    }
28}
29
30llm = OpenAI()
31chain = create_extraction_chain(schema, llm)
32result = chain.run(document_text)

Neo4j Integration

Setup

1from langchain.graphs import Neo4jGraph
2
3graph = Neo4jGraph(
4    url="bolt://localhost:7687",
5    username="neo4j",
6    password="password"
7)

Store Entities and Relationships

 1# Create nodes
 2for entity in entities:
 3    query = """
 4    MERGE (e:Entity {name: $name, type: $type})
 5    SET e.embedding = $embedding
 6    """
 7    graph.query(query, {
 8        "name": entity["name"],
 9        "type": entity["type"],
10        "embedding": get_embedding(entity["name"])
11    })
12
13# Create relationships
14for rel in relationships:
15    query = """
16    MATCH (a:Entity {name: $source})
17    MATCH (b:Entity {name: $target})
18    MERGE (a)-[r:RELATES_TO {type: $type}]->(b)
19    """
20    graph.query(query, {
21        "source": rel["source"],
22        "target": rel["target"],
23        "type": rel["type"]
24    })

Graph Retrieval Strategies

1. Direct Entity Lookup

Core Idea: Retrieves all entities directly connected to a query entity, providing immediate context about relationships.

Mathematical Formulation: $$\text{retrieve}(e_q) = {e : (e_q, r, e) \in E \lor (e, r, e_q) \in E}$$

where $e_q$ is the query entity and $E$ is the set of edges. Returns 1-hop neighbors.

Key Insight: Fast and precise for queries about direct relationships, but limited to immediate connections.

1def retrieve_entity_context(entity_name):
2    query = """
3    MATCH (e:Entity {name: $name})-[r]-(connected)
4    RETURN e, r, connected
5    LIMIT 10
6    """
7    return graph.query(query, {"name": entity_name})

2. Multi-Hop Traversal

Core Idea: Traverses the graph multiple hops from query entities, enabling reasoning across chains of relationships.

Mathematical Formulation: $$\text{retrieve}(e_q, h) = {e : \exists \text{path}(e_q, e) \text{ with length } \leq h}$$

where $h$ is the hop depth. Uses graph traversal algorithms (BFS/DFS) to find all reachable entities within $h$ hops.

Key Insight: Enables answering complex queries requiring inference across multiple relationships (e.g., "Who collaborated with authors of papers about X?").

1def retrieve_multi_hop(entity_name, hops=2):
2    query = f"""
3    MATCH path = (e:Entity {{name: $name}})-[*1..{hops}]-(connected)
4    RETURN path
5    LIMIT 20
6    """
7    return graph.query(query, {"name": entity_name})

3. Community Detection

Core Idea: Uses graph clustering algorithms (e.g., Louvain) to identify communities of related entities, then retrieves entire communities for comprehensive context.

Mathematical Formulation:

  1. Detect communities: $C = {C_1, C_2, \ldots}$ using modularity maximization: $$Q = \frac{1}{2m} \sum_{ij} \left[A_{ij} - \frac{k_i k_j}{2m}\right] \delta(c_i, c_j)$$ where $A_{ij}$ is adjacency, $k_i$ is degree, $m$ is total edges, and $\delta$ is Kronecker delta
  2. Retrieve community: $\text{retrieve}(e_q) = C_i$ where $e_q \in C_i$

Key Insight: Communities represent thematically related entities, providing broader context than direct connections alone.

 1# Run Louvain algorithm
 2graph.query("""
 3CALL gds.louvain.write({
 4    nodeProjection: 'Entity',
 5    relationshipProjection: 'RELATES_TO',
 6    writeProperty: 'community'
 7})
 8""")
 9
10# Retrieve by community
11def retrieve_community(entity_name):
12    query = """
13    MATCH (e:Entity {name: $name})
14    MATCH (related:Entity)
15    WHERE related.community = e.community
16    RETURN related
17    LIMIT 20
18    """
19    return graph.query(query, {"name": entity_name})

4. Vector + Graph Hybrid

Core Idea: Combines vector similarity search to find initial entities with graph traversal to expand context, leveraging both semantic similarity and structural relationships.

Mathematical Formulation: $$\text{retrieve}(q) = \text{expand}(\text{TopK}(\text{sim}(\mathbf{q}, \mathbf{E}), k), h)$$

where:

  1. Vector search: $\text{TopK}(\text{sim}(\mathbf{q}, \mathbf{E}), k)$ finds $k$ most similar entities
  2. Graph expansion: $\text{expand}(E_{\text{initial}}, h)$ traverses $h$ hops from initial entities
  3. Final context: Union of expanded subgraphs

Key Insight: Vector search provides semantic relevance, while graph expansion adds structural context and related entities.

 1def hybrid_retrieve(query_text, k=5):
 2    # 1. Vector search for initial entities
 3    query_embedding = get_embedding(query_text)
 4    
 5    cypher = """
 6    CALL db.index.vector.queryNodes('entity_embeddings', $k, $embedding)
 7    YIELD node, score
 8    
 9    // 2. Expand to connected entities
10    MATCH (node)-[r]-(connected)
11    RETURN node, r, connected, score
12    ORDER BY score DESC
13    """
14    
15    return graph.query(cypher, {
16        "k": k,
17        "embedding": query_embedding
18    })

LangChain Graph QA

1from langchain.chains import GraphCypherQAChain
2
3chain = GraphCypherQAChain.from_llm(
4    llm=OpenAI(),
5    graph=graph,
6    verbose=True
7)
8
9response = chain.run("Who are the authors of papers about RAG?")

Microsoft GraphRAG Approach

Core Idea: Builds a hierarchical summary structure by detecting communities, summarizing each community, and organizing summaries into a tree for multi-level querying.

Mathematical Formulation:

  1. Community Detection: $C = \text{Louvain}(G)$ partitions graph into communities
  2. Community Summarization: $S_i = \text{LLM}(\text{summarize}, C_i)$ for each community
  3. Hierarchy Construction: Build tree $T$ where nodes are communities/summaries
  4. Query Routing: $\text{query}(q) = \text{route}(q, T)$ selects appropriate level in hierarchy

Key Insight: Hierarchical organization enables querying at different granularities—detailed entity-level or high-level community summaries.

Hierarchical Summarization

 1from langchain.chains.summarize import load_summarize_chain
 2
 3# 1. Extract entities and build graph
 4# 2. Detect communities
 5# 3. Summarize each community
 6
 7def summarize_community(community_nodes):
 8    documents = [node.text for node in community_nodes]
 9    chain = load_summarize_chain(llm, chain_type="map_reduce")
10    summary = chain.run(documents)
11    return summary
12
13# 4. Build hierarchy
14# 5. Query at appropriate level

Graph RAG with LlamaIndex

 1from llama_index import KnowledgeGraphIndex, ServiceContext
 2from llama_index.graph_stores import Neo4jGraphStore
 3from llama_index.llms import OpenAI
 4
 5# Setup
 6graph_store = Neo4jGraphStore(
 7    username="neo4j",
 8    password="password",
 9    url="bolt://localhost:7687",
10    database="neo4j"
11)
12
13service_context = ServiceContext.from_defaults(
14    llm=OpenAI(model="gpt-4"),
15    chunk_size=512
16)
17
18# Build index
19index = KnowledgeGraphIndex.from_documents(
20    documents,
21    storage_context=storage_context,
22    service_context=service_context,
23    max_triplets_per_chunk=10,
24    include_embeddings=True
25)
26
27# Query
28query_engine = index.as_query_engine(
29    include_text=True,
30    response_mode="tree_summarize"
31)
32
33response = query_engine.query("What are the main concepts?")

Subgraph Extraction

Core Idea: Extracts a focused subgraph around query entities by traversing paths up to a maximum depth, creating a compact, relevant context.

Mathematical Formulation: $$\text{subgraph}(E_q, h) = {v, e : v \in V', e \in E' \land \exists \text{path}(e_q, v) \text{ with length } \leq h}$$

where $E_q$ are query entities, $h$ is max depth, and $V', E'$ form the induced subgraph.

Key Insight: Reduces noise by focusing on entities and relationships most relevant to the query, improving context quality for generation.

 1def extract_relevant_subgraph(query_entities, max_depth=2):
 2    """Extract subgraph around query entities"""
 3    
 4    query = """
 5    UNWIND $entities AS entity
 6    MATCH path = (e:Entity {name: entity})-[*1..$max_depth]-(connected)
 7    WITH collect(path) AS paths
 8    CALL apoc.convert.toTree(paths) YIELD value
 9    RETURN value
10    """
11    
12    result = graph.query(query, {
13        "entities": query_entities,
14        "max_depth": max_depth
15    })
16    
17    return result

Graph-Enhanced Prompting

 1def graph_enhanced_prompt(query, subgraph):
 2    """Create prompt with graph context"""
 3    
 4    # Convert subgraph to text
 5    context = format_subgraph_as_text(subgraph)
 6    
 7    prompt = f"""
 8    Given the following knowledge graph context:
 9    
10    {context}
11    
12    Answer the question: {query}
13    
14    Use the relationships and entities in the graph to provide a comprehensive answer.
15    """
16    
17    return llm(prompt)
18
19def format_subgraph_as_text(subgraph):
20    """Format graph as readable text"""
21    lines = []
22    for node in subgraph['nodes']:
23        lines.append(f"- {node['name']} ({node['type']})")
24    
25    for rel in subgraph['relationships']:
26        lines.append(f"  → {rel['type']} → {rel['target']}")
27    
28    return "\n".join(lines)

Evaluation

Graph Quality Metrics

 1def evaluate_graph_quality(graph):
 2    metrics = {}
 3    
 4    # Node count
 5    metrics['node_count'] = graph.query("MATCH (n) RETURN count(n) AS count")[0]['count']
 6    
 7    # Edge count
 8    metrics['edge_count'] = graph.query("MATCH ()-[r]->() RETURN count(r) AS count")[0]['count']
 9    
10    # Average degree
11    result = graph.query("""
12        MATCH (n)
13        RETURN avg(size((n)-->())) AS avg_degree
14    """)
15    metrics['avg_degree'] = result[0]['avg_degree']
16    
17    # Connected components
18    result = graph.query("""
19        CALL gds.wcc.stats({nodeProjection: '*', relationshipProjection: '*'})
20        YIELD componentCount
21        RETURN componentCount
22    """)
23    metrics['components'] = result[0]['componentCount']
24    
25    return metrics

Related Snippets