Extract entities and relations with LlamaIndex KG index

Extracting structured entities and relations from unstructured text is one of the most practical applications of LLMs, and LlamaIndex’s KnowledgeGraphIndex makes it accessible without a custom pipeline. This llamaindex entity relation extraction knowledge graph tutorial walks through building a queryable graph from raw documents, using an LLM to do the extraction work. You’ll end up with a NetworkX graph you can inspect, query, and extend.

Step 1: Install dependencies

Start with a clean environment. You need LlamaIndex core, the knowledge graph integration, and an LLM provider. We’ll use OpenAI for extraction since it handles structured output reliably, but the same pattern works with local models via Ollama or any OpenAI-compatible endpoint.

pip install llama-index llama-index-llms-openai llama-index-graph-stores-networkx networkx pyvis

If you prefer a local model, swap the LLM import and initialization in Step 3 — the rest of the pipeline stays identical.

Step 2: Load and prepare your documents

LlamaIndex’s SimpleDirectoryReader handles PDFs, Markdown, HTML, and plain text. Point it at a directory of source material. For this tutorial, create a data/ folder with a few text files — technical docs, meeting notes, or Wikipedia articles work well.

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader("data").load_data()
print(f"Loaded {len(documents)} documents")
for doc in documents[:3]:
    print(f"  - {doc.metadata.get('file_name', 'unknown')}: {len(doc.text)} chars")

Verify success: you should see a non-zero document count and character lengths that match your source files.

Step 3: Configure the LLM and extractor

The KnowledgeGraphIndex uses an LLM to extract (subject, predicate, object) triples from each text chunk. You control extraction quality through the LLM choice and the prompt template. Start with GPT-4o-mini for cost-effective extraction, or GPT-4o for higher accuracy on complex domains.

from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
import os

# Set your API key in the environment or pass directly
# os.environ["OPENAI_API_KEY"] = "sk-..."

llm = OpenAI(model="gpt-4o-mini", temperature=0.0)
Settings.llm = llm
Settings.chunk_size = 512  # Smaller chunks = more granular extraction
Settings.chunk_overlap = 50

Temperature 0.0 ensures deterministic extraction. Chunk size matters: too large and the LLM misses relations across chunk boundaries; too small and you lose context. 512 tokens is a reasonable starting point for technical text.

Step 4: Build the knowledge graph index

Now create the index. The KnowledgeGraphIndex constructor accepts the documents, an LLM (or uses Settings.llm), and optional parameters for extraction behavior.

from llama_index.core import KnowledgeGraphIndex
from llama_index.graph_stores.networkx import NetworkXGraphStore
import networkx as nx

# Optional: persist the graph to disk for reuse
graph_store = NetworkXGraphStore(nx.Graph())

index = KnowledgeGraphIndex.from_documents(
    documents,
    graph_store=graph_store,
    max_triplets_per_chunk=10,      # Limit triples per chunk to control cost
    include_embeddings=True,        # Enable hybrid vector+graph retrieval
    show_progress=True,
)

Key parameters:

• max_triplets_per_chunk: Caps extraction per chunk. Increase for dense technical text; decrease to control token spend.
• include_embeddings: When True, LlamaIndex also embeds each chunk for vector similarity search alongside graph traversal. This hybrid approach often outperforms pure graph or pure vector retrieval.

Verify success: the index builds without errors. Check the underlying graph:

print(f"Nodes: {graph_store.graph.number_of_nodes()}")
print(f"Edges: {graph_store.graph.number_of_edges()}")
# Sample a few triples
for u, v, data in list(graph_store.graph.edges(data=True))[:5]:
    print(f"  ({u}) -[{data.get('rel_type', 'rel')}]-> ({v})")

You should see dozens to hundreds of nodes and edges depending on your corpus size.

Step 5: Query the knowledge graph

LlamaIndex provides several query engines for the KG index. The KnowledgeGraphQueryEngine traverses the graph to find relevant subgraphs, then synthesizes an answer.

from llama_index.core.query_engine import KnowledgeGraphQueryEngine

query_engine = KnowledgeGraphQueryEngine(
    index=index,
    include_text=True,           # Include source text in response
    retriever_mode="keyword",    # Options: "keyword", "embedding", "hybrid"
    response_mode="tree_summarize",
    verbose=True,
)

response = query_engine.query(
    "What organizations are mentioned and what are their relationships?"
)
print(response)

Retriever modes:

• "keyword": Matches query terms to entity names in the graph. Fast, precise for known entities.
• "embedding": Uses vector similarity on chunk embeddings. Better for conceptual queries.
• "hybrid": Combines both. Recommended for production.

Verify success: the response cites specific entities and relations from your documents, not hallucinated facts. With verbose=True, you’ll see the retrieved subgraph printed to console.

Step 6: Inspect and visualize the graph

A major advantage of NetworkX storage is direct graph access. You can run graph algorithms, export to Gephi, or render in a notebook.

# Find highest-degree entities (hubs)
import networkx as nx

degrees = dict(graph_store.graph.degree())
top_entities = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:10]
print("Top entities by degree:")
for entity, deg in top_entities:
    print(f"  {entity}: {deg} connections")

# Export for Gephi
nx.write_gexf(graph_store.graph, "knowledge_graph.gexf")

For interactive browser visualization, pyvis works well:

from pyvis.network import Network

net = Network(height="750px", width="100%", bgcolor="#ffffff", font_color="black")
net.from_nx(graph_store.graph)
net.show("knowledge_graph.html", notebook=False)

Open knowledge_graph.html in a browser. You can drag nodes, zoom, and inspect edge labels (relation types).

Verify success: the HTML file opens and shows a connected graph with labeled nodes and directed edges. Hub entities appear central.

Step 7: Customize extraction with a tailored prompt

The default extraction prompt works for general domains, but specialized text (legal, biomedical, code) benefits from a custom prompt. Define a KnowledgeGraphIndex subclass or pass a custom kg_triple_extract_template.

from llama_index.core.prompts import PromptTemplate

CUSTOM_KG_TRIPLET_EXTRACT_TMPL = (
    "You are an expert in {domain} knowledge extraction.\n"
    "Extract subject-predicate-object triples from the text below.\n"
    "Focus on these relation types: {relation_types}\n"
    "Ignore generic relations like 'is', 'has', 'contains'.\n"
    "Return ONLY valid JSON array of [subject, predicate, object].\n\n"
    "Text: {text}\n\n"
    "Triples:"
)

custom_prompt = PromptTemplate(CUSTOM_KG_TRIPLET_EXTRACT_TMPL)

index_custom = KnowledgeGraphIndex.from_documents(
    documents,
    graph_store=NetworkXGraphStore(nx.Graph()),
    kg_triple_extract_template=custom_prompt.partial_format(
        domain="financial regulations",
        relation_types="regulates, acquires, subsidiaries, partners_with, filed_by"
    ),
    max_triplets_per_chunk=15,
    show_progress=True,
)

The partial_format call bakes domain context into the prompt while leaving {text} for per-chunk injection. Adjust relation_types to your ontology.

Verify success: re-run the inspection from Step 6. You should see relation types matching your specified ontology (e.g., “regulates”, “acquires”) instead of generic predicates.

Step 8: Persist and reload for production use

Rebuilding the graph on every startup is wasteful. Persist both the graph store and the vector index (if include_embeddings=True).

# Persist
index.storage_context.persist(persist_dir="./kg_storage")

# Reload later
from llama_index.core import StorageContext, load_index_from_storage

storage_context = StorageContext.from_defaults(persist_dir="./kg_storage")
index_loaded = load_index_from_storage(storage_context)

# Query engine works identically
query_engine = KnowledgeGraphQueryEngine(index=index_loaded, include_text=True)
response = query_engine.query("Summarize the regulatory relationships in the corpus.")
print(response)

The persist directory contains graph_store.json (NetworkX adjacency), docstore.json (chunks), and vector_store.json (embeddings) if enabled.

Verify success: reload completes in seconds (not minutes), and queries return identical results to the in-memory index.

Step 9: Evaluate extraction quality systematically

Extraction quality varies by domain, LLM, and prompt. Build a small evaluation set: 10-20 manually annotated triples from your documents, then measure precision/recall.

from llama_index.core.evaluation import FaithfulnessEvaluator, RelevancyEvaluator

# Faithfulness: does the answer stay grounded in retrieved triples?
faithfulness = FaithfulnessEvaluator(llm=llm)
relevancy = RelevancyEvaluator(llm=llm)

test_questions = [
    "What company acquired which subsidiary?",
    "Which regulation applies to which entity?",
]

for q in test_questions:
    response = query_engine.query(q)
    faith_result = faithfulness.evaluate_response(response=response)
    rel_result = relevancy.evaluate_response(query=q, response=response)
    print(f"Q: {q}")
    print(f"  Faithfulness: {faith_result.passing} (score: {faith_result.score})")
    print(f"  Relevancy: {rel_result.passing} (score: {rel_result.score})")

For triple-level evaluation, extract the graph’s triples and compare against your gold set:

def extract_triples(graph):
    triples = set()
    for u, v, data in graph.edges(data=True):
        rel = data.get("rel_type", "related_to")
        triples.add((u, rel, v))
    return triples

gold_triples = {
    ("Acme Corp", "acquires", "Beta Inc"),
    ("Regulation 2023/111", "applies_to", "Acme Corp"),
    # ... your annotated set
}

extracted = extract_triples(graph_store.graph)
precision = len(extracted & gold_triples) / len(extracted) if extracted else 0
recall = len(extracted & gold_triples) / len(gold_triples) if gold_triples else 0
print(f"Triple extraction - Precision: {precision:.2f}, Recall: {recall:.2f}")

Verify success: precision above 0.7 and recall above 0.5 are reasonable starting targets for first-pass extraction. Iterate on prompt and chunk size to improve.

Step 10: Extend with multi-hop reasoning

The real power of a knowledge graph is multi-hop traversal — answering questions that require chaining relations. LlamaIndex’s KnowledgeGraphQueryEngine with retriever_mode="hybrid" and explore_paths=True enables this.

query_engine_multihop = KnowledgeGraphQueryEngine(
    index=index,
    include_text=True,
    retriever_mode="hybrid",
    explore_paths=True,        # Enable multi-hop traversal
    max_paths_per_entity=3,    # Limit branching factor
    path_max_length=3,         # Max hops
    response_mode="tree_summarize",
)

response = query_engine_multihop.query(
    "Through what chain of relationships does Regulation 2023/111 affect Beta Inc?"
)
print(response)
print("\n--- Retrieved paths ---")
for path in response.metadata.get("paths", []):
    print(" -> ".join(str(node) for node in path))

With explore_paths=True, the retriever finds connecting paths between query-relevant entities. The response metadata includes the actual paths traversed.

Verify success: the answer cites a chain like “Regulation 2023/111 → applies_to → Acme Corp → acquires → Beta Inc” with each hop grounded in extracted triples.

Common pitfalls and fixes

When to use this vs. alternatives

This approach excels when:

• You need explainable, auditable extraction (every triple is traceable to source text)
• Domain relations are well-defined and you can enumerate them
• Multi-hop reasoning is required (supply chains, regulatory compliance, biography)

Consider alternatives when:

• You only need entity recognition without relations → use spaCy or GLiNER
• Relations are open-ended and you can’t predefine types → try REBEL or OpenIE
• Scale exceeds millions of documents → move to a property graph database (Neo4j, Kuzu) with LlamaIndex’s PropertyGraphIndex

The KnowledgeGraphIndex with NetworkX is a solid starting point that runs locally, persists to disk, and integrates with LlamaIndex’s broader query and evaluation tooling. For n4n.ai users routing extraction workloads across providers, the same index code works unchanged — just swap the OpenAI LLM for any OpenAI-compatible endpoint.