[Bug]: Querying neo4j knowledge graph about llama_index HOT 11 CLOSED

from llama_index.core.prompts.base import PromptTemplate, PromptType
from llama_index.core import StorageContext, ServiceContext
from llama_index.core.query_engine import KnowledgeGraphQueryEngine
from llama_index.graph_stores.neo4j import Neo4jGraphStore

# Define your custom prompt templates or use existing ones
graph_query_synthesis_prompt = PromptTemplate(
    "Your graph query synthesis prompt template here",
    prompt_type=PromptType.QUESTION_ANSWER,
)

# Initialize the LLM and ServiceContext
llm = your_llm_instance  # Replace with your actual LLM instance
service_context = ServiceContext.from_defaults(llm=llm, chunk_size=256)

# Initialize the Neo4jGraphStore
username = "neo4j"
password = "ZlXSOpUBLwCH8In2lLcPvf-beyfb87egf7e6fwefgfegf"
url = "neo4j+s://f44bb72e.databases.neo4j.io"
database = "neo4j"
graph_store = Neo4jGraphStore(
    username=username,
    password=password,
    url=url,
    database=database
)

# Create a StorageContext with your GraphStore
storage_context = StorageContext.from_defaults(graph_store=graph_store)

# Initialize the KnowledgeGraphQueryEngine with the required prompts
query_engine = KnowledgeGraphQueryEngine(
    storage_context=storage_context,
    service_context=service_context,
    llm=llm,
    graph_query_synthesis_prompt=graph_query_synthesis_prompt,
    graph_response_answer_prompt=DEFAULT_KG_RESPONSE_ANSWER_PROMPT,  # or your custom prompt
    verbose=True,
    refresh_schema=True
)

# Query the knowledge graph
response = query_engine.query(
    "Tell me more about what the author worked on at Interleaf",
)
print(response)

def run_request(self, question, history=None, retry=True):
    """
    Executes a Cypher query based on a given question.

    Args:
        question (str): The question to execute the Cypher query for.
        history (list, optional): A list of previous interactions for context. Defaults to None.
        retry (bool, optional): Whether to retry in case of a syntax error. Defaults to True.

    Returns:
        list/str: The query results or an error message.
    """
    from neo4j.exceptions import CypherSyntaxError

    # Construct Cypher statement
    cypher = self.construct_cypher_query(question, history)
    # Validate Cypher statement
    if self.cypher_query_corrector:
        cypher = self.cypher_query_corrector(cypher)
    print(cypher)
    try:
        return self.query_graph_db(cypher)
    # Self-healing flow
    except CypherSyntaxError as e:
        # If out of retries
        if not retry:
            return "Invalid Cypher syntax"
        # Self-healing Cypher flow by
        # providing specific error to GPT-4
        print("Retrying")
        return self.run_request(
            question,
            [
                ChatMessage(role=MessageRole.ASSISTANT, content=cypher),
                ChatMessage(
                    role=MessageRole.SYSTEM,
                    content=f"This query returns an error: {e!s}\n"
                    "Give me an improved query that works without any explanations or apologies",
                ),
            ],
            retry=False,
        )

class CypherQueryCorrector:
    """
    Used to correct relationship direction in generated Cypher statements.
    This code is copied from the winner's submission to the Cypher competition:
    https://github.com/sakusaku-rich/cypher-direction-competition.
    """

    # Patterns for detecting nodes, properties, and paths
    property_pattern = re.compile(r"\{.+?\}")
    node_pattern = re.compile(r"\(.+?\)")
    path_pattern = re.compile(r"\(.*\).*-.*-.*\(.*\)")
    node_relation_node_pattern = re.compile(
        r"(\()+(?P<left_node>[^()]*?)\)(?P<relation>.*?)\((?P<right_node>[^()]*?)(\))+"
    )
    relation_type_pattern = re.compile(r":(?P<relation_type>.+?)?(\{.+\})?]")

    def __init__(self, schemas: List[Schema]):
        """
        Args:
            schemas: list of schemas.
        """
        self.schemas = schemas

    def clean_node(self, node: str) -> str:
        """
        Args:
            node: node in string format.
        """
        node = re.sub(self.property_pattern, "", node)
        node = node.replace("(", "")
        node = node.replace(")", "")
        return node.strip()

    def detect_node_variables(self, query: str) -> Dict[str, List[str]]:
        """
        Args:
            query: cypher query.
        """
        nodes = re.findall(self.node_pattern, query)
        nodes = [self.clean_node(node) for node in nodes]
        res: Dict[str, Any] = {}
        for node in nodes:
            parts = node.split(":")
            if parts == "":
                continue
            variable = parts[0]
            if variable not in res:
                res[variable] = []
            res[variable] += parts[1:]
        return res

    def extract_paths(self, query: str) -> "List[str]":
        """
        Args:
            query: cypher query.
        """
        return re.findall(self.path_pattern, query)

    def judge_direction(self, relation: str) -> str:
        """
        Args:
            relation: relation in string format.
        """
        direction = "BIDIRECTIONAL"
        if relation[0] == "<":
            direction = "INCOMING"
        if relation[-1] == ">":
            direction = "OUTGOING"
        return direction

    def extract_node_variable(self, part: str) -> Optional[str]:
        """
        Args:
            part: node in string format.
        """
        part = part.lstrip("(").rstrip(")")
        idx = part.find(":")
        if idx != -1:
            part = part[:idx]
        return None if part == "" else part

    def detect_labels(
        self, str_node: str, node_variable_dict: Dict[str, Any]
    ) -> List[str]:
        """
        Args:
            str_node: node in string format
            node_variable_dict: dictionary of node variables.
        """
        splitted_node = str_node.split(":")
        variable = splitted_node[0]
        labels = []
        if variable in node_variable_dict:
            labels = node_variable_dict[variable]
        elif variable == "" and len(splitted_node) > 1:
            labels = splitted_node[1:]
        return labels

    def verify_schema(
        self,
        from_node_labels: List[str],
        relation_types: List[str],
        to_node_labels: List[str],
    ) -> bool:
        """
        Args:
            from_node_labels: labels of the from node
            relation_type: type of the relation
            to_node_labels: labels of the to node.
        """
        valid_schemas = self.schemas
        if from_node_labels != []:
            from_node_labels = [label.strip("`") for label in from_node_labels]
            valid_schemas = [
                schema for schema in valid_schemas if schema[0] in from_node_labels
            ]
        if to_node_labels != []:
            to_node_labels = [label.strip("`") for label in to_node_labels]
            valid_schemas = [
                schema for schema in valid_schemas if schema[2] in to_node_labels
            ]
        if relation_types != []:
            relation_types = [type.strip("`") for type in relation_types]
            valid_schemas = [
                schema for schema in valid_schemas if schema[1] in relation_types
            ]
        return valid_schemas != []

    def detect_relation_types(self, str_relation: str) -> Tuple[str, List[str]]:
        """
        Args:
            str_relation: relation in string format.
        """
        relation_direction = self.judge_direction(str_relation)
        relation_type = self.relation_type_pattern.search(str_relation)
        if relation_type is None or relation_type.group("relation_type") is None:
            return relation_direction, []
        relation_types = [
            t.strip().strip("!")
            for t in relation_type.group("relation_type").split("|")
        ]
        return relation_direction, relation_types

    def correct_query(self, query: str) -> str:
        """
        Args:
            query: cypher query.
        """
        node_variable_dict = self.detect_node_variables(query)
        paths = self.extract_paths(query)
        for path in paths:
            original_path = path
            start_idx = 0
            while start_idx < len(path):
                match_res = re.match(self.node_relation_node_pattern, path[start_idx:])
                if match_res is None:
                    break
                start_idx += match_res.start()
                match_dict = match_res.groupdict()
                left_node_labels = self.detect_labels(
                    match_dict["left_node"], node_variable_dict
                )
                right_node_labels = self.detect_labels(
                    match_dict["right_node"], node_variable_dict
                )
                end_idx = (
                    start_idx
                    + 4
                    + len(match_dict["left_node"])
                    + len(match_dict["relation"])
                    + len(match_dict["right_node"])
                )
                original_partial_path = original_path[start_idx : end_idx + 1]
                relation_direction, relation_types = self.detect_relation_types(
                    match_dict["relation"]
                )

                if relation_types != [] and "".join(relation_types).find("*") != -1:
                    start_idx += (
                        len(match_dict["left_node"]) + len(match_dict["relation"]) + 2
                    )
                    continue
                if relation_direction == "OUTGOING":
                    is_legal = self.verify_schema(
                        left_node_labels, relation_types, right_node_labels
                    )
                    if not is_legal:
                        is_legal = self.verify_schema(
                            right_node_labels, relation_types, left_node_labels
                        )
                        if is_legal:
                            corrected_relation = "<" + match_dict["relation"][:-1]
                            corrected_partial_path = original_partial_path.replace(
                                match_dict["relation"], corrected_relation
                            )
                            query = query.replace(
                                original_partial_path, corrected_partial_path
                            )
                        else:
                            return ""
                elif relation_direction == "INCOMING":
                    is_legal = self.verify_schema(
                        right_node_labels, relation_types, left_node_labels
                    )
                    if not is_legal:
                        is_legal = self.verify_schema(
                            left_node_labels, relation_types, right_node_labels
                        )
                        if is_legal:
                            corrected_relation = match_dict["relation"][1:] + ">"
                            corrected_partial_path = original_partial_path.replace(
                                match_dict["relation"], corrected_relation
                            )
                            query = query.replace(
                                original_partial_path, corrected_partial_path
                            )
                        else:
                            return ""
                else:
                    is_legal = self.verify_schema(
                        left_node_labels, relation_types, right_node_labels
                    )
                    is_legal |= self.verify_schema(
                        right_node_labels, relation_types, left_node_labels
                    )
                    if not is_legal:
                        return ""
                start_idx += (
                    len(match_dict["left_node"]) + len(match_dict["relation"]) + 2
                )
        return query

    def __call__(self, query: str) -> str:
        """Correct the query to make it valid. If
        Args:
            query: cypher query.
        """
        return self.correct_query(query)

from importlib.util import find_spec
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from llama_index.graph_stores.neo4j import Neo4jGraphStore
from llama_index.tools.neo4j.query_validator import CypherQueryCorrector, Schema
from neo4j.exceptions import CypherSyntaxError
from llama_index.core.llms.base import LLM, ChatMessage, MessageRole

class Neo4jQueryToolSpec(BaseToolSpec):
    """
    This class is responsible for querying a Neo4j graph database based on a provided schema definition.
    """

    spec_functions = ["run_request"]

    def __init__(self, url, user, password, database, llm: LLM, validate_cypher: bool = False):
        """
        Initializes the Neo4jSchemaWiseQuery object.

        Args:
            url (str): The connection string for the Neo4j database.
            user (str): Username for the Neo4j database.
            password (str): Password for the Neo4j database.
            llm (obj): A language model for generating Cypher queries.
            validate_cypher (bool): Validate relationship directions in the generated Cypher statement. Default: False
        """
        if find_spec("neo4j") is None:
            raise ImportError("`neo4j` package not found, please run `pip install neo4j`")

        self.graph_store = Neo4jGraphStore(url=url, username=user, password=password, database=database)
        self.llm = llm
        self.cypher_query_corrector = None
        if validate_cypher:
            corrector_schema = [
                Schema(el["start"], el["type"], el["end"])
                for el in self.graph_store.structured_schema.get("relationships")
            ]
            self.cypher_query_corrector = CypherQueryCorrector(corrector_schema)

    def get_system_message(self):
        """
        Generates a system message detailing the task and schema.
        """
        return f"""
        Task: Generate Cypher queries to query a Neo4j graph database based on the provided schema definition.
        Instructions:
        Use only the provided relationship types and properties.
        Do not use any other relationship types or properties that are not provided.
        If you cannot generate a Cypher statement based on the provided schema, explain the reason to the user.
        Schema:
        {self.graph_store.schema}

        Note: Do not include any explanations or apologies in your responses.
        """

    def query_graph_db(self, neo4j_query, params=None):
        """
        Queries the Neo4j database.

        Args:
            neo4j_query (str): The Cypher query to be executed.
            params (dict, optional): Parameters for the Cypher query. Defaults to None.

        Returns:
            list: The query results.
        """
        if params is None:
            params = {}
        with self.graph_store.client.session() as session:
            result = session.run(neo4j_query, params)
            output = [r.values() for r in result]
            output.insert(0, list(result.keys()))
            return output

    def construct_cypher_query(self, question, history=None):
        """
        Constructs a Cypher query based on a given question and history.

        Args:
            question (str): The question to construct the Cypher query for.
            history (list, optional): A list of previous interactions for context. Defaults to None.

        Returns:
            str: The constructed Cypher query.
        """
        messages = [
            ChatMessage(role=MessageRole.SYSTEM, content=self.get_system_message()),
            ChatMessage(role=MessageRole.USER, content=question),
        ]
        # Used for Cypher healing flows
        if history:
            messages.extend(history)

        completions = self.llm.chat(messages)
        return completions.message.content

    def run_request(self, question, history=None, retry=True):
        """
        Executes a Cypher query based on a given question.

        Args:
            question (str): The question to execute the Cypher query for.
            history (list, optional): A list of previous interactions for context. Defaults to None.
            retry (bool, optional): Whether to retry in case of a syntax error. Defaults to True.

        Returns:
            list/str: The query results or an error message.
        """
        # Construct Cypher statement
        cypher = self.construct_cypher_query(question, history)
        # Validate Cypher statement
        if self.cypher_query_corrector:
            cypher = self.cypher_query_corrector(cypher)
        print(cypher)
        try:
            return self.query_graph_db(cypher)
        # Self-healing flow
        except CypherSyntaxError as e:
            # If out of retries
            if not retry:
                return "Invalid Cypher syntax"
            # Self-healing Cypher flow by providing specific error to GPT-4
            print("Retrying")
            return self.run_request(
                question,
                [
                    ChatMessage(role=MessageRole.ASSISTANT, content=cypher),
                    ChatMessage(
                        role=MessageRole.SYSTEM,
                        content=f"This query returns an error: {e!s}\n"
                        "Give me an improved query that works without any explanations or apologies",
                    ),
                ],
                retry=False,
            )

from importlib.util import find_spec
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from llama_index.graph_stores.neo4j import Neo4jGraphStore
from llama_index.tools.neo4j.query_validator import CypherQueryCorrector, Schema
from neo4j.exceptions import CypherSyntaxError
from llama_index.core.llms.base import LLM, ChatMessage, MessageRole

class Neo4jQueryToolSpec(BaseToolSpec):
    spec_functions = ["run_request"]

    def __init__(self, url, user, password, database, llm: LLM, validate_cypher: bool = False):
        if find_spec("neo4j") is None:
            raise ImportError("`neo4j` package not found, please run `pip install neo4j`")

        self.graph_store = Neo4jGraphStore(url=url, username=user, password=password, database=database)
        self.llm = llm
        self.cypher_query_corrector = None
        if validate_cypher:
            corrector_schema = [
                Schema(el["start"], el["type"], el["end"])
                for el in self.graph_store.structured_schema.get("relationships")
            ]
            self.cypher_query_corrector = CypherQueryCorrector(corrector_schema)

    def get_system_message(self):
        return f"""
        Task: Generate Cypher queries to query a Neo4j graph database based on the provided schema definition.
        Instructions:
        Use only the provided relationship types and properties.
        Do not use any other relationship types or properties that are not provided.
        If you cannot generate a Cypher statement based on the provided schema, explain the reason to the user.
        Schema:
        {self.graph_store.schema}

        Note: Do not include any explanations or apologies in your responses.
        """

    def query_graph_db(self, neo4j_query, params=None):
        if params is None:
            params = {}
        with self.graph_store.client.session() as session:
            result = session.run(neo4j_query, params)
            output = [r.values() for r in result]
            output.insert(0, list(result.keys()))
            return output

    def construct_cypher_query(self, question, history=None):
        messages = [
            ChatMessage(role=MessageRole.SYSTEM, content=self.get_system_message()),
            ChatMessage(role=MessageRole.USER, content=question),
        ]
        if history:
            messages.extend(history)

        completions = self.llm.chat(messages)
        return completions.message.content

    def run_request(self, question, history=None, retry=True):
        cypher = self.construct_cypher_query(question, history)
        if self.cypher_query_corrector:
            cypher = self.cypher_query_corrector(cypher)
        print(cypher)
        try:
            return self.query_graph_db(cypher)
        except CypherSyntaxError as e:
            if not retry:
                return "Invalid Cypher syntax"
            print("Retrying")
            return self.run_request(
                question,
                [
                    ChatMessage(role=MessageRole.ASSISTANT, content=cypher),
                    ChatMessage(
                        role=MessageRole.SYSTEM,
                        content=f"This query returns an error: {e!s}\n"
                        "Give me an improved query that works without any explanations or apologies",
                    ),
                ],
                retry=False,
            )

# Initialize the Neo4jQueryToolSpec with the necessary parameters
neo4j_tool = Neo4jQueryToolSpec(
    url="bolt://localhost:7687",  # Replace with your Neo4j connection URL
    user="neo4j",                 # Replace with your Neo4j username
    password="password",          # Replace with your Neo4j password
    database="neo4j",             # Replace with your Neo4j database name
    llm=your_llm_instance,        # Replace with your LLM instance
    validate_cypher=True          # Set to True if you want to validate Cypher queries
)

# Define the question you want to query
question = "What are the names of all employees in the database?"

# Execute the query
results = neo4j_tool.run_request(question)

# Print the results
print(results)

from llama_index.core.prompts.base import PromptTemplate, PromptType
from llama_index.core import StorageContext, ServiceContext
from llama_index.core.query_engine import KnowledgeGraphQueryEngine
from llama_index.graph_stores.neo4j import Neo4jGraphStore
from llama_index.tools.neo4j.query_validator import CypherQueryCorrector, Schema
from neo4j.exceptions import CypherSyntaxError

# Define your custom prompt templates
graph_query_synthesis_prompt = PromptTemplate(
    "Generate a Cypher query to retrieve information based on the following question: {query_str}",
    prompt_type=PromptType.QUESTION_ANSWER,
)

DEFAULT_KG_RESPONSE_ANSWER_PROMPT = PromptTemplate(
    "Provide a detailed and comprehensive answer to the query, including any relevant context and background information.",
    prompt_type=PromptType.QUESTION_ANSWER,
)

# Initialize the LLM and ServiceContext
llm = your_llm_instance  # Replace with your actual LLM instance
service_context = ServiceContext.from_defaults(llm=llm, chunk_size=256)

# Initialize the Neo4jGraphStore
username = "neo4j"
password = "your_password"  # Replace with your Neo4j password
url = "your_neo4j_url"  # Replace with your Neo4j URL
database = "neo4j"

graph_store = Neo4jGraphStore(
    username=username,
    password=password,
    url=url,
    database=database
)

# Create a StorageContext with your GraphStore
storage_context = StorageContext.from_defaults(graph_store=graph_store)

# Initialize the KnowledgeGraphQueryEngine with the required prompts
query_engine = KnowledgeGraphQueryEngine(
    storage_context=storage_context,
    service_context=service_context,
    llm=llm,
    graph_query_synthesis_prompt=graph_query_synthesis_prompt,
    graph_response_answer_prompt=DEFAULT_KG_RESPONSE_ANSWER_PROMPT,  # or your custom prompt
    verbose=True,
    refresh_schema=True
)

# Define the CypherQueryCorrector
schemas = [
    Schema("Person", "KNOWS", "Person"),
    Schema("Person", "LIKES", "Movie"),
    # Add more schemas as needed
]
cypher_query_corrector = CypherQueryCorrector(schemas)

# Function to run the query with self-healing flow
def run_request(query_engine, question, history=None, retry=True):
    cypher = query_engine.generate_query(question)
    cypher = cypher_query_corrector(cypher)
    print("Generated Cypher Query:", cypher)
    try:
        return query_engine.query_graph_db(cypher)
    except CypherSyntaxError as e:
        if not retry:
            return "Invalid Cypher syntax"
        print("Retrying with corrected query")
        return run_request(
            query_engine,
            question,
            [
                ChatMessage(role=MessageRole.ASSISTANT, content=cypher),
                ChatMessage(
                    role=MessageRole.SYSTEM,
                    content=f"This query returns an error: {e!s}\n"
                    "Give me an improved query that works without any explanations or apologies",
                ),
            ],
            retry=False,
        )

# Query the knowledge graph
response = run_request(query_engine, "tell me what government did to prevent it")
print(response)