A modular Graph RAG (Retrieval-Augmented Generation) library using Qdrant and OpenAI. Built for a CS workshop on building knowledge graph-powered chatbots from scratch.
GraphRAG combines vector similarity search with graph traversal to provide rich, contextual responses. The library features:
- Plug-and-play schemas: Generic, Knowledge Graph, and Document-centric schemas
- LLM-based extraction: Extract entities and relationships from raw text
- Multi-hop retrieval: Traverse graph relationships to find connected context
- Real-time visualization: See the graph "think" as queries are processed
- Modular architecture: Swap components easily for experimentation
┌─────────────────────────────────────────────────────────────────┐
│ GraphRAG │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │
│ │ Schemas │ │Extraction│ │ Storage │ │ Retrieval │ │
│ │ (3 types)│ │ (LLM) │ │ (Qdrant) │ │ (Multi-hop) │ │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └────────┬─────────┘ │
│ │ │ │ │ │
│ └─────────────┴─────────────┴──────────────────┘ │
│ │ │
│ ┌────────────────────────┴────────────────────────────────────┐ │
│ │ Chat Agent │ │
│ │ (RAG with Graph Excitation) │ │
│ └──────────────────────────────────────────────────────────────┘ │
│ │ │
│ ┌────────────────────────┴────────────────────────────────────┐ │
│ │ Visualization │ │
│ │ (Plotly 3D / React WebApp) │ │
│ └──────────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
GraphRAG/
├── graphrag/ # Core library
│ ├── config.py # Configuration management
│ ├── embeddings/ # Embedding models (OpenAI)
│ ├── storage/ # Qdrant storage layer
│ ├── schema/ # Graph schemas (generic, knowledge, document)
│ ├── extraction/ # Entity/relationship extraction
│ ├── construction/ # Graph building pipeline
│ ├── retrieval/ # Multi-hop traversal & context
│ ├── chat/ # RAG chatbot agent
│ ├── io/ # Data loaders & exporters
│ └── visualization/ # Plotly 3D & server
├── webapp/ # React 3D visualization app
├── data/ # Sample datasets
├── demo/ # Demo scripts (01-05)
└── workshop/ # Student assignments
- Python 3.10+ with Poetry
- Qdrant (local or cloud)
- OpenAI API key
# Clone the repository
git clone <repo-url>
cd GraphRAG
# Install dependencies with Poetry
poetry install
# Or with pip
pip install openai qdrant-client plotly python-dotenvOption A: Docker (Local)
docker run -p 6333:6333 -p 6334:6334 qdrant/qdrantOption B: Qdrant Cloud
- Create account at cloud.qdrant.io
- Create a cluster and get your URL + API key
Create a .env file:
# Required
OPENAI_API_KEY=sk-your-api-key
# For Qdrant Cloud (optional)
QDRANT_URL=https://your-cluster.qdrant.io
QDRANT_API_KEY=your-qdrant-api-key
# For local Qdrant (default)
QDRANT_HOST=localhost
QDRANT_PORT=6333# 1. Basic setup and connection test
python demo/01_basic_setup.py
# 2. Build a graph from sample data
python demo/02_construction.py
# 3. Query the graph
python demo/03_retrieval.py
# 4. Visualize in 3D
python demo/04_visualization.py
# 5. Interactive chatbot
python demo/05_chatbot.pyGraphRAG uses two Qdrant collections:
-
Nodes Collection: Each point represents a graph node
- Vector: Embedding of node text
- Payload:
{text, node_type, edges: [{target_id, edge_type, ...}], properties}
-
Edges Collection: Each point represents an edge (for reverse lookups)
- Vector: Embedding of relationship description
- Payload:
{source_id, target_id, edge_type, properties}
Query: "How does deep learning relate to AI?"
Step 1 (Vector Search):
→ Find nodes similar to query
→ Result: ["Deep Learning", "Neural Networks"]
Step 2 (Graph Traversal - Hop 1):
→ Follow edges from seed nodes
→ Discover: ["Machine Learning", "CNN", "RNN"]
Step 3 (Graph Traversal - Hop 2):
→ Follow edges again
→ Discover: ["Artificial Intelligence", "Data Science"]
Result: Subgraph with activation scores showing "thinking" pathThree built-in schemas:
- Generic: Flexible
Entity → RELATES_TO → Entity - Knowledge Graph: Typed entities (Person, Org, Location, Event, Concept)
- Document: Document-centric (Document, Chunk, Entity, Topic)
Create custom schemas by extending GraphSchema:
from graphrag.schema.base import GraphSchema, NodeType, EdgeType, SchemaRegistry
@SchemaRegistry.register
class MySchema(GraphSchema):
@property
def name(self) -> str:
return "my_schema"
@property
def node_types(self) -> Dict[str, NodeType]:
return {"MyEntity": NodeType(...)}
@property
def edge_types(self) -> Dict[str, EdgeType]:
return {"MY_RELATION": EdgeType(...)}from graphrag.config import GraphRAGConfig
from graphrag.construction.pipeline import GraphConstructionPipeline
config = GraphRAGConfig.from_env()
with GraphConstructionPipeline(config, schema_name="knowledge_graph") as pipeline:
pipeline.setup(recreate=True)
documents = [
"OpenAI developed GPT-4. Sam Altman is the CEO of OpenAI.",
"Google owns DeepMind. DeepMind created AlphaGo."
]
stats = pipeline.process_documents(documents)
print(f"Created {stats.nodes_created} nodes, {stats.edges_created} edges")from graphrag.retrieval.traversal import GraphTraverser
with GraphTraverser(config) as traverser:
# Simple search
results = traverser.search_nodes("artificial intelligence", top_k=5)
# Multi-hop traversal
traversal = traverser.traverse("Who leads AI companies?")
# Find related entities
related = traverser.find_related("OpenAI", max_hops=2)from graphrag.chat.agent import GraphRAGAgent
with GraphRAGAgent(config) as agent:
response = agent.chat("What is the relationship between GPT-4 and OpenAI?")
print(response.message.content)
print(f"Used {len(response.context.nodes)} nodes from graph")from graphrag.visualization.plotly_viz import create_3d_graph
fig = create_3d_graph(
nodes=response.context.nodes,
edges=response.context.relationships,
activation_scores=response.context.activation_scores,
title="Query Results"
)
fig.show()For real-time graph visualization with the chatbot:
cd webapp
npm install
npm run devThen start the Python server:
python -m graphrag.visualization.serverOpen http://localhost:3000 to see the interactive visualization.
The workshop/ folder contains assignments for learning GraphRAG:
| Assignment | Topic | Skills |
|---|---|---|
| 01 | Setup & Configuration | Qdrant connection, collections |
| 02 | Custom Schema | Node/edge types, validation |
| 03 | Custom Extractor | Rule-based & LLM extraction |
| 04 | Advanced Retrieval | 3-hop traversal, weighted search |
Solutions are in workshop/solutions/.
GraphRAGConfig(
qdrant=QdrantConfig(
url="...", # For cloud
api_key="...", # For cloud
host="localhost", # For local
port=6333, # For local
nodes_collection="nodes",
edges_collection="edges"
),
openai=OpenAIConfig(
api_key="...",
embedding_model="text-embedding-3-small",
chat_model="gpt-4o-mini"
),
default_top_k=5,
max_hops=2
)| Class | Description |
|---|---|
GraphRAGConfig |
Central configuration |
GraphConstructionPipeline |
Build graphs from documents |
GraphTraverser |
Query and traverse graphs |
GraphRAGAgent |
Chat with graph context |
GraphVisualizer |
3D Plotly visualization |
LLMExtractor |
GPT-based entity extraction |
MIT
Contributions welcome! Please read the contributing guidelines first.