Knowledge Graph RAG: Hybrid Architecture for Complex Enterprise Reasoning

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.graphs import Neo4jGraph
from langchain_community.vectorstores import Neo4jVector
from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_cohere import CohereRerank
from langchain.retrievers import ContextualCompressionRetriever
from langchain.schema import Document
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# ---- Configuration ----
NEO4J_URI = "neo4j+s://your-aura-db.databases.neo4j.io"
NEO4J_USERNAME = "neo4j"
NEO4J_PASSWORD = "your-password"

llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")


class HybridRAGPipeline:
    """
    Hybrid retrieval pipeline combining:
    1. Neo4j vector store for semantic similarity retrieval
    2. Neo4j graph traversal for relationship queries via Cypher
    3. Cohere reranker to merge and rank results from both sources
    """

    def __init__(self):
        # Vector store: documents embedded into Neo4j with vector index
        self.vector_store = Neo4jVector.from_existing_index(
            embedding=embeddings,
            url=NEO4J_URI,
            username=NEO4J_USERNAME,
            password=NEO4J_PASSWORD,
            index_name="document_embeddings",
            node_label="Document",
            text_node_property="content",
            embedding_node_property="embedding"
        )

        # Graph connection for Cypher queries
        self.graph = Neo4jGraph(
            url=NEO4J_URI,
            username=NEO4J_USERNAME,
            password=NEO4J_PASSWORD
        )

        # Graph QA chain: LLM generates Cypher from natural language, executes against Neo4j
        self.cypher_chain = GraphCypherQAChain.from_llm(
            llm=llm,
            graph=self.graph,
            verbose=False,
            return_intermediate_steps=True,
            cypher_prompt=self._get_cypher_generation_prompt()
        )

        # Reranker for merging results from both retrieval paths
        self.reranker = CohereRerank(model="rerank-english-v3.0", top_n=6)
        self.compression_retriever = ContextualCompressionRetriever(
            base_compressor=self.reranker,
            base_retriever=self.vector_store.as_retriever(search_kwargs={"k": 15})
        )

    def _get_cypher_generation_prompt(self) -> ChatPromptTemplate:
        """Prompt that instructs the LLM to generate safe, read-only Cypher queries."""
        return ChatPromptTemplate.from_messages([
            ("system",
             """You are an expert Neo4j Cypher query generator for an enterprise knowledge graph.

Graph schema:
- Nodes: (Document {{id, title, content, source}}), (Entity {{id, name, type, description}}),
  (Regulation {{id, name, jurisdiction, effective_date}}), (Product {{id, name, category}}),
  (Organisation {{id, name, type}}), (Contract {{id, title, value, signed_date}})
- Relationships: (Entity)-[:MENTIONED_IN]->(Document), (Regulation)-[:APPLIES_TO]->(Product),
  (Organisation)-[:PARTY_TO]->(Contract), (Organisation)-[:SUPPLIES]->(Product),
  (Entity)-[:RELATED_TO]->(Entity)

Rules:
- Generate READ-ONLY Cypher (MATCH/RETURN only, never MERGE/CREATE/DELETE)
- Always LIMIT results to maximum 20 unless explicitly asked for more
- Return only relevant properties, not entire nodes
- If the question cannot be answered via graph traversal, return an empty Cypher query"""),
            ("human", "Schema: {schema}\n\nQuestion: {question}\n\nCypher query:")
        ])

    def _run_graph_query(self, query: str) -> list[Document]:
        """Run the graph QA chain and convert results to Document objects."""
        try:
            result = self.cypher_chain.invoke({"query": query})
            answer = result.get("result", "")
            if answer and answer.strip() and "don't know" not in answer.lower():
                return [Document(
                    page_content=f"Graph Query Result: {answer}",
                    metadata={"source": "neo4j_graph", "query_type": "cypher"}
                )]
        except Exception as e:
            logger.warning(f"graph_query_failed | error={str(e)}")
        return []

    def _run_vector_search(self, query: str, k: int = 10) -> list[Document]:
        """Run vector similarity search."""
        try:
            return self.compression_retriever.invoke(query)
        except Exception as e:
            logger.warning(f"vector_search_failed | error={str(e)}")
            return []

    def retrieve(self, query: str) -> list[Document]:
        """
        Hybrid retrieval: run both graph and vector search,
        deduplicate, and return merged results.
        """
        graph_docs = self._run_graph_query(query)
        vector_docs = self._run_vector_search(query)

        # Merge: graph results first (they are relationship-specific),
        # then vector results, deduplicating by content hash
        seen_content = set()
        merged = []
        for doc in graph_docs + vector_docs:
            content_hash = hash(doc.page_content[:200])
            if content_hash not in seen_content:
                seen_content.add(content_hash)
                merged.append(doc)

        logger.info(f"hybrid_retrieval | query_len={len(query)} | graph_docs={len(graph_docs)} | vector_docs={len(vector_docs)} | merged={len(merged)}")
        return merged[:8]  # Cap final context at 8 documents

    def query(self, question: str) -> str:
        """Full hybrid RAG query: retrieve from both sources, generate answer."""
        docs = self.retrieve(question)

        context = "\n\n".join(
            f"[Source: {d.metadata.get('source', 'unknown')}]\n{d.page_content}"
            for d in docs
        )

        if not context.strip():
            return "I could not find relevant information in the knowledge base to answer this question."

        prompt = ChatPromptTemplate.from_template("""
You are an enterprise knowledge assistant with access to both a document corpus
and a structured knowledge graph. Use the provided context to answer the question.
Cite the source type (document or graph) for key facts.

Context:
{context}

Question: {question}

Answer:
""")

        chain = prompt | llm | StrOutputParser()
        return chain.invoke({"context": context, "question": question})


if __name__ == "__main__":
    pipeline = HybridRAGPipeline()

    # Pure vector query — answered by document chunks
    answer1 = pipeline.query("What are the payment terms in our standard MSA template?")
    print("Vector query answer:", answer1[:200])

    # Multi-hop graph query — answered by relationship traversal
    answer2 = pipeline.query("Which of our products are subject to GDPR and also have contracts expiring in 2025?")
    print("Graph query answer:", answer2[:200])

    # Hybrid query — requires both sources
    answer3 = pipeline.query("What does our contract with Meridian Analytics say about liability, and are they also our competitor's supplier?")
    print("Hybrid query answer:", answer3[:200])