Neptune graph and openCypher QA Chain

docs/extras/modules/chains/additional/neptune_cypher_qa.ipynb

@@ -0,0 +1,52 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Neptune Open Cypher QA Chain\n",
+    "This QA chain queries Neptune graph database using openCypher and returns human readable response\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.graphs.neptune_graph import NeptuneGraph\n",
+    "\n",
+    "\n",
+    "host = \"<neptune-host>\"\n",
+    "port = 80\n",
+    "use_https = False\n",
+    "\n",
+    "graph = NeptuneGraph(host=host, port=port, use_https=use_https)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.chains.graph_qa.neptune_cypher import NeptuneOpenCypherQAChain\n",
+    "\n",
+    "llm = ChatOpenAI(temperature=0, model=\"gpt-4\")\n",
+    "\n",
+    "chain = NeptuneOpenCypherQAChain.from_llm(llm=llm, graph=graph)\n",
+    "\n",
+    "chain.run(\"how many outgoing routes does the Austin airport have?\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

langchain/chains/graph_qa/neptune_cypher.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import re
+from typing import Any, Dict, List, Optional
+
+from pydantic import Field
+
+from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import CallbackManagerForChainRun
+from langchain.chains.base import Chain
+from langchain.chains.graph_qa.prompts import (
+    CYPHER_QA_PROMPT,
+    NEPTUNE_OPENCYPHER_GENERATION_PROMPT,
+)
+from langchain.chains.llm import LLMChain
+from langchain.graphs import NeptuneGraph
+from langchain.prompts.base import BasePromptTemplate
+
+INTERMEDIATE_STEPS_KEY = "intermediate_steps"
+
+
+def extract_cypher(text: str) -> str:
+    # The pattern to find Cypher code enclosed in triple backticks
+    pattern = r"```(.*?)```"
+
+    # Find all matches in the input text
+    matches = re.findall(pattern, text, re.DOTALL)
+
+    return matches[0] if matches else text
+
+
+class NeptuneOpenCypherQAChain(Chain):
+    """Chain for question-answering against a Neptune graph
+    by generating openCypher statements.
+
+    Example:
+        .. code-block:: python
+
+        chain = NeptuneOpenCypherQAChain.from_llm(
+            llm=llm,
+            graph=graph
+        )
+        response = chain.run(query)
+    """
+
+    graph: NeptuneGraph = Field(exclude=True)
+    cypher_generation_chain: LLMChain
+    qa_chain: LLMChain
+    input_key: str = "query"  #: :meta private:
+    output_key: str = "result"  #: :meta private:
+    top_k: int = 10
+    return_intermediate_steps: bool = False
+    """Whether or not to return the intermediate steps along with the final answer."""
+    return_direct: bool = False
+    """Whether or not to return the result of querying the graph directly."""
+
+    @property
+    def input_keys(self) -> List[str]:
+        """Return the input keys.
+
+        :meta private:
+        """
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> List[str]:
+        """Return the output keys.
+
+        :meta private:
+        """
+        _output_keys = [self.output_key]
+        return _output_keys
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLanguageModel,
+        *,
+        qa_prompt: BasePromptTemplate = CYPHER_QA_PROMPT,
+        cypher_prompt: BasePromptTemplate = NEPTUNE_OPENCYPHER_GENERATION_PROMPT,
+        **kwargs: Any,
+    ) -> NeptuneOpenCypherQAChain:
+        """Initialize from LLM."""
+        qa_chain = LLMChain(llm=llm, prompt=qa_prompt)
+        cypher_generation_chain = LLMChain(llm=llm, prompt=cypher_prompt)
+
+        return cls(
+            qa_chain=qa_chain,
+            cypher_generation_chain=cypher_generation_chain,
+            **kwargs,
+        )
+
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        """Generate Cypher statement, use it to look up in db and answer question."""
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        callbacks = _run_manager.get_child()
+        question = inputs[self.input_key]
+
+        intermediate_steps: List = []
+
+        generated_cypher = self.cypher_generation_chain.run(
+            {"question": question, "schema": self.graph.get_schema}, callbacks=callbacks
+        )
+
+        # Extract Cypher code if it is wrapped in backticks
+        generated_cypher = extract_cypher(generated_cypher)
+
+        _run_manager.on_text("Generated Cypher:", end="\n", verbose=self.verbose)
+        _run_manager.on_text(
+            generated_cypher, color="green", end="\n", verbose=self.verbose
+        )
+
+        intermediate_steps.append({"query": generated_cypher})
+
+        context = self.graph.query(generated_cypher)
+
+        if self.return_direct:
+            final_result = context
+        else:
+            _run_manager.on_text("Full Context:", end="\n", verbose=self.verbose)
+            _run_manager.on_text(
+                str(context), color="green", end="\n", verbose=self.verbose
+            )
+
+            intermediate_steps.append({"context": context})
+
+            result = self.qa_chain(
+                {"question": question, "context": context},
+                callbacks=callbacks,
+            )
+            final_result = result[self.qa_chain.output_key]
+
+        chain_result: Dict[str, Any] = {self.output_key: final_result}
+        if self.return_intermediate_steps:
+            chain_result[INTERMEDIATE_STEPS_KEY] = intermediate_steps
+
+        return chain_result

langchain/chains/graph_qa/prompts.py

@@ -196,3 +196,21 @@
 SPARQL_QA_PROMPT = PromptTemplate(
     input_variables=["context", "prompt"], template=SPARQL_QA_TEMPLATE
 )
+
+
+NEPTUNE_OPENCYPHER_EXTRA_INSTRUCTIONS = """
+Instructions:
+Generate the query in openCypher format and follow these rules:
+Do not use `NONE`, `ALL` or `ANY` predicate functions, rather use list comprehensions.
+Do not use `REDUCE` function. Rather use a combination of list comprehension and the `UNWIND` clause to achieve similar results.
+Do not use `FOREACH` clause. Rather use a combination of `WITH` and `UNWIND` clauses to achieve similar results.
+\n"""
+
+NEPTUNE_OPENCYPHER_GENERATION_TEMPLATE = CYPHER_GENERATION_TEMPLATE.replace(
+    "Instructions:", NEPTUNE_OPENCYPHER_EXTRA_INSTRUCTIONS
+)
+
+NEPTUNE_OPENCYPHER_GENERATION_PROMPT = PromptTemplate(
+    input_variables=["schema", "question"],
+    template=NEPTUNE_OPENCYPHER_GENERATION_TEMPLATE,
+)

langchain/graphs/__init__.py

@@ -3,13 +3,15 @@
 from langchain.graphs.kuzu_graph import KuzuGraph
 from langchain.graphs.nebula_graph import NebulaGraph
 from langchain.graphs.neo4j_graph import Neo4jGraph
+from langchain.graphs.neptune_graph import NeptuneGraph
 from langchain.graphs.networkx_graph import NetworkxEntityGraph
 from langchain.graphs.rdf_graph import RdfGraph
 
 __all__ = [
     "NetworkxEntityGraph",
     "Neo4jGraph",
     "NebulaGraph",
+    "NeptuneGraph",
     "KuzuGraph",
     "HugeGraph",
     "RdfGraph",