GraphRAG Complete Guide 2026: Knowledge Graph Enhanced RAG

# Requirements: python>=3.10
# pip install langchain langchain-community langchain-openai neo4j python-dotenv networkx

import os
from dotenv import load_dotenv

load_dotenv()

# Configure OpenAI
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "")

from typing import List
from pydantic import BaseModel, Field

class Entity(BaseModel):
    """Represents an entity in the knowledge graph."""
    name: str = Field(description="The unique name of the entity")
    entity_type: str = Field(description="The type of entity (e.g., Person, Organization, Technology)")
    description: str = Field(description="A brief description of the entity")
    properties: dict = Field(default_factory=dict, description="Additional properties of the entity")

class Relationship(BaseModel):
    """Represents a relationship between two entities."""
    source: str = Field(description="The name of the source entity")
    target: str = Field(description="The name of the target entity")
    relationship_type: str = Field(description="The type of relationship (e.g., founded, works_at)")
    properties: dict = Field(default_factory=dict, description="Relationship properties")

class ExtractedGraph(BaseModel):
    """Container for extracted entities and relationships."""
    entities: List[Entity] = Field(default_factory=list)
    relationships: List[Relationship] = Field(default_factory=list)

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser

# Initialize the LLM
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# Create the extraction prompt
EXTRACTION_PROMPT = """You are an expert at extracting structured information from text.

Extract entities and relationships from the following text. Be thorough and extract all relevant entities and their relationships.

Schema:
- Entity types: Person, Organization, Technology, Product, Concept, Location, Event
- Relationship types: founded, created, developed, works_at, acquired, invested_in, uses, related_to

Text:
{text}

{format_instructions}
"""

# Create parser
parser = PydanticOutputParser(pydantic_object=ExtractedGraph)

# Create prompt template
prompt = ChatPromptTemplate.from_template(EXTRACTION_PROMPT)
prompt = prompt.partial(format_instructions=parser.get_format_instructions())

# Create extraction chain
extraction_chain = prompt | llm.with_structured_output(ExtractedGraph)

from langchain.text_splitter import RecursiveCharacterTextSplitter
from typing import List, Dict, Any
import json

class DocumentProcessor:
    """Processes documents to extract knowledge graph data."""
    
    def __init__(self, extraction_chain, chunk_size: int = 1000, chunk_overlap: int = 200):
        self.extraction_chain = extraction_chain
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", ". ", " "]
        )
    
    def process_document(self, text: str, metadata: Dict[str, Any] = None) -> ExtractedGraph:
        """Process a single document and extract graph data."""
        chunks = self.text_splitter.split_text(text)
        
        all_entities = []
        all_relationships = []
        
        # Process each chunk
        for chunk in chunks:
            try:
                result = self.extraction_chain.invoke({"text": chunk})
                all_entities.extend(result.entities)
                all_relationships.extend(result.relationships)
            except Exception as e:
                print(f"Error processing chunk: {e}")
                continue
        
        # Deduplicate entities
        entities_by_name = {}
        for entity in all_entities:
            if entity.name not in entities_by_name:
                entities_by_name[entity.name] = entity
            else:
                # Merge descriptions if entity already exists
                existing = entities_by_name[entity.name]
                if len(entity.description) > len(existing.description):
                    existing.description = entity.description
                # Merge properties
                existing.properties.update(entity.properties)
        
        return ExtractedGraph(
            entities=list(entities_by_name.values()),
            relationships=all_relationships
        )
    
    def process_documents(self, documents: List[Dict]) -> List[ExtractedGraph]:
        """Process multiple documents."""
        results = []
        for doc in documents:
            text = doc.get("text", "")
            metadata = doc.get("metadata", {})
            result = self.process_document(text, metadata)
            results.append(result)
        return results

from neo4j import GraphDatabase

class KnowledgeGraph:
    """Manages the Neo4j knowledge graph."""
    
    def __init__(self, uri: str, user: str, password: str):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
    
    def close(self):
        self.driver.close()
    
    def create_constraints(self):
        """Create constraints for unique entity names."""
        with self.driver.session() as session:
            session.run("""
                CREATE CONSTRAINT entity_name_unique 
                IF NOT EXISTS
                FOR (e:Entity) REQUIRE e.name IS UNIQUE
            """)
    
    def clear_graph(self):
        """Clear all nodes and relationships."""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
    
    def insert_extracted_graph(self, graph: ExtractedGraph):
        """Insert entities and relationships into the graph."""
        with self.driver.session() as session:
            # Insert entities
            for entity in graph.entities:
                session.run("""
                    MERGE (e:Entity {name: $name})
                    SET e.type = $type,
                        e.description = $description,
                        e.properties = $properties
                """, 
                    name=entity.name,
                    type=entity.entity_type,
                    description=entity.description,
                    properties=json.dumps(entity.properties)
                )
            
            # Insert relationships
            for rel in graph.relationships:
                session.run("""
                    MATCH (s:Entity {name: $source})
                    MATCH (t:Entity {name: $target})
                    MERGE (s)-[r:RELATES_TO {type: $rel_type}]->(t)
                    SET r.properties = $properties
                """,
                    source=rel.source,
                    target=rel.target,
                    rel_type=rel.relationship_type,
                    properties=json.dumps(rel.properties)
                )
    
    def get_entity_info(self, entity_name: str) -> Dict:
        """Get information about a specific entity."""
        with self.driver.session() as session:
            result = session.run("""
                MATCH (e:Entity {name: $name})
                OPTIONAL MATCH (e)-[r]-(related)
                RETURN e, collect(DISTINCT {relationship: type(r), entity: related.name}) as related
            """, name=entity_name)
            
            record = result.single()
            if record:
                return {
                    "entity": dict(record["e"]),
                    "related": record["related"]
                }
            return None
    
    def find_paths(self, start_entity: str, end_entity: str, max_depth: int = 3) -> List:
        """Find paths between two entities."""
        with self.driver.session() as session:
            result = session.run("""
                MATCH path = (start:Entity {name: $start})-[:RELATES_TO*1..%d]-(end:Entity {name: $end})
                RETURN path
                LIMIT 10
            """ % max_depth, start=start_entity, end=end_entity)
            
            paths = []
            for record in result:
                path = record["path"]
                nodes = [node.get("name") for node in path.nodes]
                rels = [rel.get("type") for rel in path.relationships]
                paths.append({
                    "nodes": nodes,
                    "relationships": rels
                })
            return paths

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Neo4jVector
from typing import List, Dict, Any

class GraphRAGRetriever:
    """Retrieves information from the knowledge graph."""
    
    def __init__(self, kg: KnowledgeGraph, embeddings: OpenAIEmbeddings = None):
        self.kg = kg
        self.embeddings = embeddings or OpenAIEmbeddings()
    
    def retrieve_by_entity(self, query: str, entity_name: str, depth: int = 2) -> Dict:
        """Retrieve information starting from a specific entity."""
        with self.kg.driver.session() as session:
            result = session.run("""
                MATCH (e:Entity {name: $name})-[r:RELATES_TO*1..%d]-(related)
                RETURN e, collect(DISTINCT related) as related_entities
            """ % depth, name=entity_name)
            
            record = result.single()
            if record:
                return {
                    "central_entity": dict(record["e"]),
                    "related": [dict(n) for n in record["related_entities"]]
                }
            return None
    
    def retrieve_subgraph(self, query: str, top_k: int = 10) -> Dict:
        """Retrieve a relevant subgraph based on the query."""
        # First, find entities that might be relevant using keyword matching
        keywords = query.lower().split()
        
        with self.kg.driver.session() as session:
            # Build a query to find matching entities
            placeholders = ", ".join([f"${kw}" for kw in keywords])
            result = session.run(f"""
                MATCH (e:Entity)
                WHERE ANY(keyword IN [{placeholders}] WHERE toLower(e.name) CONTAINS keyword 
                       OR toLower(e.description) CONTAINS keyword)
                OPTIONAL MATCH (e)-[r]-(related)
                RETURN e, collect(DISTINCT {{entity: related.name, relationship: type(r)}}) as related
                LIMIT $limit
            """, **{kw: kw for kw in keywords}, limit=top_k)
            
            entities = []
            for record in result:
                entities.append({
                    "entity": dict(record["e"]),
                    "related": record["related"]
                })
            
            return {"entities": entities}
    
    def hybrid_retrieve(self, query: str, vector_store, top_k: int = 5, graph_depth: int = 2) -> Dict:
        """Combine vector search with graph traversal."""
        # Vector search
        vector_results = vector_store.similarity_search(query, k=top_k)
        
        # Extract entity names from vector results (simplified)
        entity_names = []
        for doc in vector_results:
            # In practice, you'd use NER or entity linking here
            text = doc.page_content
            words = text.split()
            entity_names.extend([w for w in words if w[0].isupper() and len(w) > 3][:3])
        
        # Graph retrieval for found entities
        graph_context = []
        for name in set(entity_names[:5]):
            entity_info = self.kg.get_entity_info(name)
            if entity_info:
                graph_context.append(entity_info)
        
        return {
            "vector_results": vector_results,
            "graph_context": graph_context
        }

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda

class GraphRAGGenerator:
    """Generates answers using retrieved graph context."""
    
    def __init__(self, llm, retriever):
        self.llm = llm
        self.retriever = retriever
    
    def generate_answer(self, query: str, context: Dict) -> str:
        """Generate an answer using graph context."""
        
        # Format graph context
        graph_context = self._format_context(context)
        
        prompt = ChatPromptTemplate.from_template("""You are a helpful AI assistant with access to a knowledge graph.

Use the following graph context to answer the question. The graph contains entities and their relationships.

Graph Context:
{graph_context}

Question: {question}

Instructions:
- Use only information from the provided graph context
- If the graph doesn't contain enough information, say so
- Cite the entities and relationships you used
- Be precise and accurate
""")
        
        chain = prompt | self.llm
        
        response = chain.invoke({
            "graph_context": graph_context,
            "question": query
        })
        
        return response.content
    
    def _format_context(self, context: Dict) -> str:
        """Format graph context as text."""
        lines = []
        
        if "entities" in context:
            lines.append("## Entities and Their Relationships\n")
            for item in context["entities"]:
                entity = item.get("entity", {})
                lines.append(f"### {entity.get('name', 'Unknown')} ({entity.get('type', 'Unknown')})")
                lines.append(f"Description: {entity.get('description', 'N/A')}")
                
                related = item.get("related", [])
                if related:
                    lines.append("Related entities:")
                    for rel in related:
                        if isinstance(rel, dict):
                            ent = rel.get("entity", {})
                            rel_type = rel.get("relationship", "related_to")
                            lines.append(f"  - {ent.get('name', 'Unknown')} ({rel_type})")
                lines.append("")
        
        return "\n".join(lines)
    
    def query(self, query: str, vector_store = None) -> str:
        """Full query pipeline: retrieve and generate."""
        # Retrieve relevant context
        if vector_store:
            context = self.retriever.hybrid_retrieve(query, vector_store)
        else:
            context = self.retriever.retrieve_subgraph(query)
        
        # Generate answer
        answer = self.generate_answer(query, context)
        
        return answer

# Initialize components
kg = KnowledgeGraph(
    uri=os.getenv("NEO4J_URI", "bolt://localhost:7687"),
    user=os.getenv("NEO4J_USER", "neo4j"),
    password=os.getenv("NEO4J_PASSWORD", "password")
)

processor = DocumentProcessor(extraction_chain)
retriever = GraphRAGRetriever(kg)
generator = GraphRAGGenerator(llm, retriever)

# Process sample documents
sample_documents = [
    {
        "text": """
        OpenAI was founded in 2015 by Sam Altman, Elon Musk, and other prominent technologists. 
        The company is based in San Francisco and has become a leader in artificial intelligence research. 
        OpenAI developed GPT-3, GPT-4, and ChatGPT. Elon Musk also founded Tesla, SpaceX, and Neuralink.
        Sam Altman previously served as president of Y Combinator.
        """,
        "metadata": {"source": "tech_companies.txt"}
    },
    {
        "text": """
        Tesla was founded by Elon Musk, JB Straubel, and Martin Eberhard in 2003. 
        The company specializes in electric vehicles and clean energy. 
        Elon Musk has been instrumental in Tesla's growth and became CEO.
        SpaceX, also founded by Elon Musk, focuses on space exploration.
        """,
        "metadata": {"source": "tesla_info.txt"}
    }
]

# Extract and store knowledge graph
for doc in sample_documents:
    graph = processor.process_document(doc["text"], doc["metadata"])
    kg.insert_extracted_graph(graph)

# Query the knowledge graph
answer = generator.query("What companies did Elon Musk found?")
print(answer)

# Clean up
kg.close()

from langchain_community.graphs import Neo4jGraph
from langchain.chains import GraphQAChain
from langchain_openai import ChatOpenAI

# Initialize Neo4j graph
graph = Neo4jGraph(
    url="bolt://localhost:7687",
    username="neo4j",
    password="password"
)

# Refresh schema
graph.refresh_schema()

# Create QA chain
llm = ChatOpenAI(model="gpt-4o")
chain = GraphQAChain.from_llm(llm, graph=graph, verbose=True)

# Query
chain.run("What companies were founded by Elon Musk?")

class HybridRAGSystem:
    """Combines vector, graph, and keyword search."""
    
    def __init__(self, vector_store, kg: KnowledgeGraph, embeddings):
        self.vector_store = vector_store
        self.kg = kg
        self.embeddings = embeddings
    
    def retrieve(self, query: str, weights: Dict[str, float] = None) -> List[Dict]:
        """Retrieve from multiple sources and combine results."""
        weights = weights or {"vector": 0.4, "graph": 0.4, "keyword": 0.2}
        
        # Vector search
        vector_results = self.vector_store.similarity_search(query, k=10)
        
        # Graph search (simplified)
        graph_results = self._graph_search(query)
        
        # Keyword search
        keyword_results = self._keyword_search(query)
        
        # Combine using weighted reciprocal rank fusion
        combined = self._reciprocal_rank_fusion(
            vector_results, graph_results, keyword_results, weights
        )
        
        return combined[:10]
    
    def _graph_search(self, query: str) -> List[Dict]:
        # Implementation of graph-based search
        pass
    
    def _keyword_search(self, query: str) -> List[Dict]:
        # Implementation of keyword search
        pass
    
    def _reciprocal_rank_fusion(self, *result_sets, weights: Dict) -> List:
        """Combine ranked lists using weighted RRF."""
        scores = {}
        
        for source_name, results in zip(["vector", "graph", "keyword"], result_sets):
            weight = weights.get(source_name, 1.0)
            for rank, result in enumerate(results):
                key = self._get_result_key(result)
                scores[key] = scores.get(key, 0) + weight / (rank + 1)
        
        sorted_results = sorted(scores.items(), key=lambda x: x[1], reverse=True)
        return [item[0] for item in sorted_results]
    
    def _get_result_key(self, result):
        return str(result)  # Simplified - use actual identifier in practice

class MultiModalGraphRAG:
    """Handles text, images, and structured data in the same graph."""
    
    def __init__(self, kg: KnowledgeGraph):
        self.kg = kg
        self.image_entities = {}  # Map image IDs to extracted information
    
    def process_image(self, image_path: str, description: str, related_entities: List[str]):
        """Process an image and link it to existing entities."""
        import uuid
        image_id = str(uuid.uuid4())
        
        # Store image information
        self.image_entities[image_id] = {
            "path": image_path,
            "description": description
        }
        
        # Link to related entities in the graph
        with self.kg.driver.session() as session:
            for entity_name in related_entities:
                session.run("""
                    MATCH (e:Entity {name: $entity})
                    MERGE (img:Image {id: $id})
                    SET img.description = $desc, img.path = $path
                    MERGE (e)-[:depicted_in]->(img)
                """, entity=entity_name, id=image_id, desc=description, path=image_path)
    
    def query_with_images(self, query: str) -> Dict:
        """Query the graph and include relevant images."""
        # First retrieve relevant text entities
        text_results = self.kg.retrieve_subgraph(query)
        
        # Then find linked images
        image_results = []
        for entity in text_results.get("entities", []):
            entity_name = entity.get("name")
            with self.kg.driver.session() as session:
                result = session.run("""
                    MATCH (e:Entity {name: $name})-[:depicted_in]->(img:Image)
                    RETURN img
                """, name=entity_name)
                
                for record in result:
                    image_results.append(dict(record["img"]))
        
        return {
            "text_entities": text_results,
            "images": image_results
        }

class ExperimentTracker:
    """Track A/B test results for different retrieval strategies."""
    
    def __init__(self):
        self.experiments = {}
    
    def record_query(self, query: str, strategy: str, result: Dict):
        """Record a query result for analysis."""
        if strategy not in self.experiments:
            self.experiments[strategy] = []
        self.experiments[strategy].append({
            "query": query,
            "result": result,
            "timestamp": datetime.now()
        })
    
    def compare_strategies(self, metric_fn):
        """Compare strategies using a custom metric function."""
        results = {}
        for strategy, queries in self.experiments.items():
            scores = [metric_fn(q["result"]) for q in queries]
            results[strategy] = {
                "mean": sum(scores) / len(scores),
                "count": len(scores)
            }
        return results