Fix missing inclusion of neo4j_schema in custom prompt generation (#256)

* Fix missing inclusion of neo4j_schema in custom prompt generation

* Add documentation for schema and examples in prompt_params

Author: koyonkym

src/neo4j_graphrag/retrievers/text2cypher.py

@@ -55,7 +55,7 @@ class Text2CypherRetriever(Retriever):
         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 not include the neo4j_schema or examples args, if provided.
+        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.
@@ -99,7 +99,13 @@ def __init__(
         self.result_formatter = validated_data.result_formatter
         self.custom_prompt = validated_data.custom_prompt
         if validated_data.custom_prompt:
-            neo4j_schema = ""
+            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
@@ -124,7 +130,7 @@ def get_search_results(
 
         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. Example: {'schema': 'this is the graph schema'}
+            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.

tests/unit/retrievers/test_text2cypher.py

@@ -254,6 +254,73 @@ def test_t2c_retriever_initialization_with_custom_prompt_and_schema_and_examples
     llm.invoke.assert_called_once_with("This is a custom prompt. test")
 
 
+@patch("neo4j_graphrag.retrievers.base.get_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_schema_and_examples_for_prompt_params(
+    mock_get_version: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    mock_get_version.return_value = ((5, 23, 0), False, False)
+    prompt = "This is a custom prompt. {query_text} {schema} {examples}"
+    neo4j_schema = "dummy-schema"
+    examples = ["example-1", "example-2"]
+
+    retriever = Text2CypherRetriever(
+        driver=driver,
+        llm=llm,
+        custom_prompt=prompt,
+        neo4j_schema=neo4j_schema,
+        examples=examples,
+    )
+
+    driver.execute_query.return_value = (
+        [neo4j_record],
+        None,
+        None,
+    )
+    retriever.search(query_text="test")
+
+    llm.invoke.assert_called_once_with(
+        "This is a custom prompt. test dummy-schema example-1\nexample-2"
+    )
+
+
+@patch("neo4j_graphrag.retrievers.base.get_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_unused_schema_and_examples(
+    mock_get_version: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    mock_get_version.return_value = ((5, 23, 0), False, False)
+    prompt = "This is a custom prompt. {query_text} {schema} {examples}"
+    neo4j_schema = "dummy-schema"
+    examples = ["example-1", "example-2"]
+
+    retriever = Text2CypherRetriever(
+        driver=driver,
+        llm=llm,
+        custom_prompt=prompt,
+        neo4j_schema=neo4j_schema,
+        examples=examples,
+    )
+
+    driver.execute_query.return_value = (
+        [neo4j_record],
+        None,
+        None,
+    )
+    retriever.search(
+        query_text="test",
+        prompt_params={"schema": "another-dummy-schema", "examples": "another-example"},
+    )
+
+    llm.invoke.assert_called_once_with(
+        "This is a custom prompt. test another-dummy-schema another-example"
+    )
+
+
 @patch("neo4j_graphrag.retrievers.base.get_version")
 def test_t2c_retriever_invalid_custom_prompt_type(
     mock_get_version: MagicMock, driver: MagicMock, llm: MagicMock