Fix missing inclusion of neo4j_schema in custom prompt generation

Author: koyonkym

src/neo4j_graphrag/retrievers/text2cypher.py

@@ -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

tests/unit/retrievers/test_text2cypher.py

@@ -245,6 +245,71 @@ 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.Retriever._verify_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_schema_and_examples_for_prompt_params(
+    _verify_version_mock: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    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.Retriever._verify_version")
+def test_t2c_retriever_initialization_with_custom_prompt_and_unused_schema_and_examples(
+    _verify_version_mock: MagicMock,
+    driver: MagicMock,
+    llm: MagicMock,
+    neo4j_record: MagicMock,
+) -> None:
+    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.Text2CypherRetriever._verify_version")
 def test_t2c_retriever_invalid_custom_prompt_type(
     _verify_version_mock: MagicMock, driver: MagicMock, llm: MagicMock