RAG with Knowledge Graphs: Neo4j Integration Guide

Layer 1: Vector RAG (semantic similarity)
         ↓
Layer 2: Entity Extraction (NER, LLM)
         ↓
Layer 3: Knowledge Graph (Neo4j relationships)

Documents → Extract Entities & Relations → Neo4j Graph
                                            ↓
Query → Classify → Vector Search | Cypher Query | Hybrid
                         ↓              ↓            ↓
                    Chunks        Graph Data    Combined
                         ↓              ↓            ↓
                         → Generate Answer ←

from openai import OpenAI
from neo4j import AsyncGraphDatabase
from typing import List, Dict

client = OpenAI()

async def extract_entities_and_relationships(document: str) -> dict:
    """Extract entities and relationships using LLM."""

    extract_prompt = f"""Extract entities and relationships from the document.

Document:
{document}

Return JSON:
{{
  "entities": [
    {{"name": str, "type": str, "properties": dict}}
  ],
  "relationships": [
    {{"source": str, "target": str, "type": str, "properties": dict}}
  ]
}}

Entity types: Person, Organization, Service, API, Database, Document, Concept
Relationship types: DEPENDS_ON, CONNECTS_TO, PART_OF, REFERENCES, AUTHORED_BY, MANAGES
"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": extract_prompt}],
        response_format={"type": "json_object"},
        temperature=0.1
    )

    return json.loads(response.choices[0].message.content)
# {
#   "entities": [
#     {"name": "AuthService", "type": "Service", "properties": {"language": "Python"}},
#     {"name": "UserDB", "type": "Database", "properties": {"type": "PostgreSQL"}},
#     {"name": "PaymentAPI", "type": "API", "properties": {"version": "v2"}}
#   ],
#   "relationships": [
#     {"source": "AuthService", "target": "UserDB", "type": "CONNECTS_TO", "properties": {}},
#     {"source": "PaymentAPI", "target": "AuthService", "type": "DEPENDS_ON", "properties": {"critical": true}}
#   ]
# }

class Neo4jGraphBuilder:
    """Build knowledge graph from extracted entities."""

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

    async def create_entities(self, entities: List[Dict]):
        """Create entity nodes in Neo4j."""

        async with self.driver.session() as session:
            for entity in entities:
                await session.run(
                    f"""
                    MERGE (n:{entity['type']} {{name: $name}})
                    SET n += $properties
                    """,
                    name=entity["name"],
                    properties=entity.get("properties", {})
                )

    async def create_relationships(self, relationships: List[Dict]):
        """Create relationship edges in Neo4j."""

        async with self.driver.session() as session:
            for rel in relationships:
                await session.run(
                    f"""
                    MATCH (a {{name: $source}}), (b {{name: $target}})
                    MERGE (a)-[r:{rel['type']}]->(b)
                    SET r += $properties
                    """,
                    source=rel["source"],
                    target=rel["target"],
                    properties=rel.get("properties", {})
                )

    async def ingest_document(self, document: str, metadata: Dict):
        """Full pipeline: extract entities, build graph."""

        # Step 1: Extract
        extracted = await extract_entities_and_relationships(document)

        # Step 2: Create nodes
        await self.create_entities(extracted["entities"])

        # Step 3: Create edges
        await self.create_relationships(extracted["relationships"])

        # Step 4: Link document to entities
        await self._link_document_to_entities(metadata["source"], extracted["entities"])

    async def _link_document_to_entities(self, doc_id: str, entities: List[Dict]):
        """Link document node to mentioned entities."""

        async with self.driver.session() as session:
            # Create document node
            await session.run(
                """
                MERGE (d:Document {id: $doc_id})
                """,
                doc_id=doc_id
            )

            # Link to entities
            for entity in entities:
                await session.run(
                    """
                    MATCH (d:Document {id: $doc_id}), (e {name: $entity_name})
                    MERGE (d)-[:MENTIONS]->(e)
                    """,
                    doc_id=doc_id,
                    entity_name=entity["name"]
                )

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

# Usage
graph_builder = Neo4jGraphBuilder(
    uri="bolt://localhost:7687",
    user="neo4j",
    password="password"
)

document = """
The AuthService connects to UserDB for authentication.
PaymentAPI depends on AuthService for authorization.
"""

await graph_builder.ingest_document(document, {"source": "architecture-doc.md"})

async def text_to_cypher(question: str, schema: str) -> str:
    """Generate Cypher query from natural language."""

    cypher_prompt = f"""You are a Neo4j Cypher expert. Generate a Cypher query to answer the question.

Graph Schema:
{schema}

Question: {question}

Return JSON: {{"cypher": str, "reasoning": str}}

Examples:
Q: "What services depend on AuthService?"
A: MATCH (s:Service)-[:DEPENDS_ON]->(target:Service {{name: 'AuthService'}}) RETURN s.name

Q: "How are ServiceA and ServiceB connected?"
A: MATCH path = shortestPath((a:Service {{name: 'ServiceA'}})-[*]-(b:Service {{name: 'ServiceB'}})) RETURN path
"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": cypher_prompt}],
        response_format={"type": "json_object"},
        temperature=0.1
    )

    result = json.loads(response.choices[0].message.content)
    return result["cypher"]

# Usage
schema = """
Nodes:
- Service (name, language, version)
- Database (name, type)
- API (name, version)

Relationships:
- DEPENDS_ON (from Service to Service/Database/API)
- CONNECTS_TO (from Service to Database)
- REFERENCES (from Document to Entity)
"""

cypher = await text_to_cypher("What databases does PaymentService connect to?", schema)
# Result: MATCH (s:Service {name: 'PaymentService'})-[:CONNECTS_TO]->(d:Database) RETURN d.name

async def execute_cypher(cypher: str, driver) -> List[Dict]:
    """Execute Cypher query and return results."""

    async with driver.session() as session:
        result = await session.run(cypher)
        records = await result.data()
        return records

async def graph_rag_query(question: str, graph_builder: Neo4jGraphBuilder) -> str:
    """Full text-to-Cypher RAG pipeline."""

    # Step 1: Get graph schema
    schema = await get_graph_schema(graph_builder.driver)

    # Step 2: Generate Cypher
    cypher = await text_to_cypher(question, schema)

    # Step 3: Execute query
    results = await execute_cypher(cypher, graph_builder.driver)

    # Step 4: Format answer
    answer = await format_graph_results(question, results)

    return answer

async def format_graph_results(question: str, results: List[Dict]) -> str:
    """Convert graph query results to natural language."""

    results_text = json.dumps(results, indent=2)

    format_prompt = f"""Format the graph query results as a natural language answer.

Question: {question}

Results:
{results_text}

Provide a clear, concise answer.
"""

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": format_prompt}],
        temperature=0.1
    )

    return response.choices[0].message.content

class HybridGraphRAG:
    """Hybrid RAG: vector search + graph traversal."""

    def __init__(self, vector_conn, graph_builder: Neo4jGraphBuilder):
        self.vector_conn = vector_conn
        self.graph_builder = graph_builder

    async def query(self, question: str) -> dict:
        """Hybrid query: vector + graph."""

        # Step 1: Classify query type
        classification = await self._classify_query(question)

        vector_results = []
        graph_results = []

        # Step 2: Vector search (if needed)
        if classification["needs_vector"]:
            vector_results = await self._vector_search(question)

        # Step 3: Graph traversal (if needed)
        if classification["needs_graph"]:
            graph_results = await self._graph_search(question)

        # Step 4: Synthesize answer from both
        answer = await self._synthesize(question, vector_results, graph_results)

        return {
            "answer": answer,
            "vector_sources": vector_results,
            "graph_data": graph_results,
        }

    async def _classify_query(self, question: str) -> dict:
        """Classify if query needs vector search, graph search, or both."""

        classify_prompt = f"""Classify the query:

Question: {question}

Return JSON:
{{
  "needs_vector": bool,
  "needs_graph": bool,
  "reasoning": str
}}

needs_vector: True if query asks about document content, concepts, definitions
needs_graph: True if query asks about relationships, dependencies, connections
"""

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": classify_prompt}],
            response_format={"type": "json_object"}
        )

        return json.loads(response.choices[0].message.content)

    async def _vector_search(self, question: str) -> List[str]:
        """Retrieve relevant chunks via vector search."""
        from pgvector.asyncpg import register_vector

        query_embedding = client.embeddings.create(
            input=[question],
            model="text-embedding-3-small"
        ).data[0].embedding

        await register_vector(self.vector_conn)
        chunks = await self.vector_conn.fetch(
            """
            SELECT content
            FROM document_chunks
            ORDER BY embedding <=> $1
            LIMIT 5
            """,
            query_embedding
        )

        return [c["content"] for c in chunks]

    async def _graph_search(self, question: str) -> List[Dict]:
        """Retrieve via graph traversal."""
        schema = await get_graph_schema(self.graph_builder.driver)
        cypher = await text_to_cypher(question, schema)
        return await execute_cypher(cypher, self.graph_builder.driver)

    async def _synthesize(self, question: str, vector_results: List[str], graph_results: List[Dict]) -> str:
        """Synthesize answer from vector + graph results."""

        context_parts = []

        if vector_results:
            context_parts.append("Document Context:\n" + "\n\n---\n\n".join(vector_results))

        if graph_results:
            context_parts.append(f"Graph Data:\n{json.dumps(graph_results, indent=2)}")

        combined_context = "\n\n".join(context_parts)

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "Answer using document and graph data."},
                {"role": "user", "content": f"{combined_context}\n\nQuestion: {question}"}
            ],
            temperature=0.1
        )

        return response.choices[0].message.content

# Example queries
# "What is Kubernetes?" → Vector search (concept definition)
# "What services depend on AuthService?" → Graph traversal (relationships)
# "How does the auth flow work and what services are involved?" → Hybrid (concept + dependencies)

from dataclasses import dataclass

@dataclass
class GraphRAGConfig:
    neo4j_uri: str
    neo4j_user: str
    neo4j_password: str
    enable_vector: bool = True
    enable_graph: bool = True
    auto_extract_entities: bool = True

class ProductionGraphRAG:
    """Production RAG with Neo4j knowledge graph."""

    def __init__(self, config: GraphRAGConfig, vector_conn):
        self.config = config
        self.vector_conn = vector_conn
        self.graph_builder = Neo4jGraphBuilder(
            uri=config.neo4j_uri,
            user=config.neo4j_user,
            password=config.neo4j_password
        )
        self.hybrid_rag = HybridGraphRAG(vector_conn, self.graph_builder)

    async def ingest_document(self, document: str, metadata: Dict):
        """Ingest document into both vector and graph stores."""

        # Step 1: Traditional vector RAG ingestion
        if self.config.enable_vector:
            await self._ingest_vector(document, metadata)

        # Step 2: Extract entities and build graph
        if self.config.enable_graph and self.config.auto_extract_entities:
            await self.graph_builder.ingest_document(document, metadata)

    async def query(self, question: str) -> dict:
        """Query using hybrid vector + graph approach."""

        if self.config.enable_vector and self.config.enable_graph:
            return await self.hybrid_rag.query(question)

        elif self.config.enable_graph:
            # Graph-only
            graph_results = await self.hybrid_rag._graph_search(question)
            answer = await format_graph_results(question, graph_results)
            return {"answer": answer, "graph_data": graph_results}

        else:
            # Vector-only
            vector_results = await self.hybrid_rag._vector_search(question)
            context = "\n\n---\n\n".join(vector_results)
            answer = await self._generate_from_context(question, context)
            return {"answer": answer, "vector_sources": vector_results}

    async def visualize_graph(self, entity: str) -> str:
        """Generate visualization of entity and relationships."""

        cypher = f"""
        MATCH (e {{name: $entity}})-[r]-(connected)
        RETURN e, r, connected
        LIMIT 20
        """

        results = await execute_cypher(cypher, self.graph_builder.driver)

        # Return as DOT format for Graphviz
        dot = "digraph G {\n"
        for record in results:
            src = record["e"]["name"]
            rel = type(record["r"]).__name__
            dst = record["connected"]["name"]
            dot += f'  "{src}" -> "{dst}" [label="{rel}"];\n'
        dot += "}"

        return dot

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

# 1. Index frequently queried properties
async def create_indexes(driver):
    async with driver.session() as session:
        await session.run("CREATE INDEX service_name IF NOT EXISTS FOR (s:Service) ON (s.name)")
        await session.run("CREATE INDEX api_name IF NOT EXISTS FOR (a:API) ON (a.name)")

# 2. Limit traversal depth
cypher = """
MATCH path = (a:Service {name: 'AuthService'})-[*1..3]-(b)
RETURN path
LIMIT 50
"""

# 3. Use query profiling
cypher_with_profile = """
PROFILE
MATCH (s:Service)-[:DEPENDS_ON]->(d)
RETURN s, d
"""

# 4. Cache common queries
from functools import lru_cache

@lru_cache(maxsize=1000)
async def get_service_dependencies(service_name: str) -> List[str]:
    # Expensive graph query
    pass