# Copyright (c) "Neo4j" # Neo4j Sweden AB [https://neo4j.com] # # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # # https://www.apache.org/licenses/LICENSE-2.0 # # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import logging import re from typing import Any, Callable, Dict, Optional import neo4j from neo4j.exceptions import CypherSyntaxError, DriverError, Neo4jError from pydantic import ValidationError from neo4j_graphrag.exceptions import ( RetrieverInitializationError, SchemaFetchError, SearchValidationError, Text2CypherRetrievalError, ) from neo4j_graphrag.generation.prompts import Text2CypherTemplate from neo4j_graphrag.llm import LLMInterface from neo4j_graphrag.retrievers.base import Retriever from neo4j_graphrag.schema import get_schema from neo4j_graphrag.types import ( LLMModel, Neo4jDriverModel, Neo4jSchemaModel, RawSearchResult, RetrieverResultItem, Text2CypherRetrieverModel, Text2CypherSearchModel, ) logger = logging.getLogger(__name__) def extract_cypher(text: str) -> str: """Extract and format Cypher query from text, handling code blocks and special characters. This function performs two main operations: 1. Extracts Cypher code from within triple backticks (```), if present 2. Automatically adds backtick quotes around multi-word identifiers: - Node labels (e.g., ":Data Science" becomes ":`Data Science`") - Property keys (e.g., "first name:" becomes "`first name`:") - Relationship types (e.g., "[:WORKS WITH]" becomes "[:`WORKS WITH`]") Args: text (str): Raw text that may contain Cypher code, either within triple backticks or as plain text. Returns: str: Properly formatted Cypher query with correct backtick quoting. """ # Extract Cypher code enclosed in triple backticks pattern = r"```(.*?)```" matches = re.findall(pattern, text, re.DOTALL) cypher_query = matches[0] if matches else text # Quote node labels in backticks if they contain spaces and are not already quoted cypher_query = re.sub( r":\s*(?!`\s*)(\s*)([a-zA-Z0-9_]+(?:\s+[a-zA-Z0-9_]+)+)(?!\s*`)(\s*)", r":`\2`", cypher_query, ) # Quote property keys in backticks if they contain spaces and are not already quoted cypher_query = re.sub( r"([,{]\s*)(?!`)([a-zA-Z0-9_]+(?:\s+[a-zA-Z0-9_]+)+)(?!`)(\s*:)", r"\1`\2`\3", cypher_query, ) # Quote relationship types in backticks if they contain spaces and are not already quoted cypher_query = re.sub( r"(\[\s*[a-zA-Z0-9_]*\s*:\s*)(?!`)([a-zA-Z0-9_]+(?:\s+[a-zA-Z0-9_]+)+)(?!`)(\s*(?:\]|-))", r"\1`\2`\3", cypher_query, ) return cypher_query class Text2CypherRetriever(Retriever): """ Allows for the retrieval of records from a Neo4j database using natural language. Converts a user's natural language query to a Cypher query using an LLM, then retrieves records from a Neo4j database using the generated Cypher query. Args: driver (neo4j.Driver): The Neo4j Python driver. llm (neo4j_graphrag.generation.llm.LLMInterface): LLM object to generate the Cypher query. neo4j_schema (Optional[str]): Neo4j schema used to generate the Cypher query. examples (Optional[list[str], optional): Optional user input/query pairs for the LLM to use as examples. custom_prompt (Optional[str]): Optional custom prompt to use instead of auto generated prompt. Will include the neo4j_schema for schema and examples for examples prompt parameters, if they are provided. Raises: RetrieverInitializationError: If validation of the input arguments fail. """ def __init__( self, driver: neo4j.Driver, llm: LLMInterface, neo4j_schema: Optional[str] = None, examples: Optional[list[str]] = None, result_formatter: Optional[ Callable[[neo4j.Record], RetrieverResultItem] ] = None, custom_prompt: Optional[str] = None, neo4j_database: Optional[str] = None, ) -> None: try: driver_model = Neo4jDriverModel(driver=driver) llm_model = LLMModel(llm=llm) neo4j_schema_model = ( Neo4jSchemaModel(neo4j_schema=neo4j_schema) if neo4j_schema else None ) validated_data = Text2CypherRetrieverModel( driver_model=driver_model, llm_model=llm_model, neo4j_schema_model=neo4j_schema_model, examples=examples, result_formatter=result_formatter, custom_prompt=custom_prompt, neo4j_database=neo4j_database, ) except ValidationError as e: raise RetrieverInitializationError(e.errors()) from e super().__init__( validated_data.driver_model.driver, validated_data.neo4j_database ) self.llm = validated_data.llm_model.llm self.examples = validated_data.examples self.result_formatter = validated_data.result_formatter self.custom_prompt = validated_data.custom_prompt if validated_data.custom_prompt: if ( validated_data.neo4j_schema_model and validated_data.neo4j_schema_model.neo4j_schema ): neo4j_schema = validated_data.neo4j_schema_model.neo4j_schema else: neo4j_schema = "" else: if ( validated_data.neo4j_schema_model and validated_data.neo4j_schema_model.neo4j_schema ): neo4j_schema = validated_data.neo4j_schema_model.neo4j_schema else: try: neo4j_schema = get_schema(validated_data.driver_model.driver) except (Neo4jError, DriverError) as e: error_message = getattr(e, "message", str(e)) raise SchemaFetchError( f"Failed to fetch schema for Text2CypherRetriever: {error_message}" ) from e self.neo4j_schema = neo4j_schema def get_search_results( self, query_text: str, prompt_params: Optional[Dict[str, Any]] = None ) -> RawSearchResult: """Converts query_text to a Cypher query using an LLM. Retrieve records from a Neo4j database using the generated Cypher query. Args: query_text (str): The natural language query used to search the Neo4j database. prompt_params (Dict[str, Any]): additional values to inject into the custom prompt, if it is provided. If the schema or examples parameter is specified, it will overwrite the corresponding value passed during initialization. Example: {'schema': 'this is the graph schema'} Raises: SearchValidationError: If validation of the input arguments fail. Text2CypherRetrievalError: If the LLM fails to generate a correct Cypher query. Returns: RawSearchResult: The results of the search query as a list of neo4j.Record and an optional metadata dict """ try: validated_data = Text2CypherSearchModel(query_text=query_text) except ValidationError as e: raise SearchValidationError(e.errors()) from e prompt_template = Text2CypherTemplate(template=self.custom_prompt) if prompt_params is not None: # parse the schema and examples inputs examples_to_use = prompt_params.pop("examples", None) or ( "\n".join(self.examples) if self.examples else "" ) schema_to_use = prompt_params.pop("schema", None) or self.neo4j_schema else: examples_to_use = "\n".join(self.examples) if self.examples else "" schema_to_use = self.neo4j_schema prompt_params = dict() prompt = prompt_template.format( schema=schema_to_use, examples=examples_to_use, query_text=validated_data.query_text, **prompt_params, ) logger.debug("Text2CypherRetriever prompt: %s", prompt) try: llm_result = self.llm.invoke(prompt) t2c_query = extract_cypher(llm_result.content) logger.debug("Text2CypherRetriever Cypher query: %s", t2c_query) records, _, _ = self.driver.execute_query( query_=t2c_query, database_=self.neo4j_database, routing_=neo4j.RoutingControl.READ, ) except CypherSyntaxError as e: raise Text2CypherRetrievalError( f"Failed to get search result: {e.message}" ) from e return RawSearchResult( records=records, metadata={ "cypher": t2c_query, }, )