Architecting High-Performance Knowledge Graphs: Integrating LLMs with Neo4j and Async Python

Beyond Vector Search: The Rise of GraphRAG

While Vector Databases have revolutionized semantic search, they often struggle with complex, multi-hop reasoning. A vector search can find documents similar to a query, but it fails to understand the structural relationships between entities—such as why a specific CEO is linked to a subsidiary's regulatory failure across different continents. This is where Knowledge Graphs (KGs) come in. By representing data as a network of nodes and relationships, we provide LLMs with a structured world model.

In this post, we will explore how to architect a high-throughput pipeline that extracts structured triplets (Subject-Predicate-Object) from unstructured text using LLMs and persists them into Neo4j using asynchronous Python. This architecture, often referred to as GraphRAG, enables more sophisticated querying than traditional RAG systems.

The Architectural Blueprint

To build a production-grade Knowledge Graph pipeline, we need to balance the latency of LLM inference with the transactional integrity of graph database writes. Our stack includes:

1. Extraction Layer: Using an LLM (like GPT-4o or Claude 3.5) to parse raw text into JSON-formatted triplets.
2. Orchestration Layer: Async Python (asyncio) to handle concurrent LLM calls and database I/O.
3. Storage Layer: Neo4j, utilizing the native Python driver’s asynchronous capabilities.

Designing the Extraction Logic

The first challenge is ensuring the LLM outputs consistent schema. We define a clear prompt that forces the model to identify entities and their relationships.

import asyncio
from neo4j import AsyncGraphDatabase
from typing import List, Dict
import json

# Example schema for the extracted triplets
TRIPLET_PROMPT = """
Extract entities and their relationships from the following text.
Return a JSON list of objects with 'subject', 'predicate', and 'object'.
Text: {text}
"""

async def extract_triplets(text: str, llm_client) -> List[Dict[str, str]]:
    # In a real scenario, call your LLM provider here
    response = await llm_client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": TRIPLET_PROMPT.format(text=text)}],
        response_format={ "type": "json_object" }
    )
    return json.loads(response.choices[0].message.content).get("triplets", [])

High-Performance Graph Ingestion

Writing to Neo4j one node at a time is a performance killer. To achieve high throughput, we must leverage Cypher's UNWIND clause to batch operations and use the AsyncGraphDatabase driver to prevent blocking the event loop.

Async Neo4j Driver Implementation

class KnowledgeGraphStore:
    def __init__(self, uri, user, password):
        self.driver = AsyncGraphDatabase.driver(uri, auth=(user, password))

    async def close(self):
        await self.driver.close()

    async def write_triplets_batch(self, triplets: List[Dict[str, str]]):
        query = """
        UNWIND $batch AS item
        MERGE (s:Entity {name: item.subject})
        MERGE (o:Entity {name: item.object})
        WITH s, o, item
        CALL apoc.create.relationship(s, item.predicate, {}, o) YIELD rel
        RETURN count(rel)
        """
        async with self.driver.session() as session:
            await session.execute_write(self._execute_query, query, triplets)

    @staticmethod
    async self._execute_query(tx, query, batch):
        await tx.run(query, batch=batch)

Optimizing the Pipeline for Scale

1. Connection Pooling and Concurrency Control

When dealing with thousands of documents, you cannot simply spawn a task for every document. Use a semaphore to limit the number of concurrent LLM requests and database sessions to avoid exhausting file descriptors or hitting rate limits.

async def process_documents(docs: List[str], kg_store: KnowledgeGraphStore):
    semaphore = asyncio.Semaphore(10)  # Limit to 10 concurrent extractions

    async def worker(doc):
        async with semaphore:
            triplets = await extract_triplets(doc)
            await kg_store.write_triplets_batch(triplets)

    await asyncio.gather(*(worker(doc) for doc in docs))

2. Entity Disambiguation

One of the hardest problems in Knowledge Graphs is ensuring that "Apple Inc." and "Apple" refer to the same node. Before writing to the graph, implementing a light-weight entity resolution step—either via fuzzy matching or a second LLM pass—is crucial for graph density and accuracy.

3. Cypher Query Optimization

Ensure that you have indexes on the properties used in MERGE statements. In Neo4j, running CREATE INDEX FOR (e:Entity) ON (e.name) is mandatory for performance. Without it, every insertion becomes a full-label scan, leading to exponential slowdowns as the graph grows.

Conclusion: The Future of Context-Aware AI

Building a Knowledge Graph is more labor-intensive than setting up a simple vector index, but the dividends are immense. By structuring your data as a graph, you enable your AI agents to perform complex reasoning, maintain long-term memory, and provide explainable insights. As the industry moves toward "Agentic Workflows," the ability to orchestrate these high-performance graph pipelines will become a defining skill for AI engineers.

In the next phase of this architecture, one could integrate Graph Neural Networks (GNNs) to predict missing links within the KG, further enriching the context available to the LLM.