Fix bug in Text2cypher custom prompt option (#115)

a-s-g93 · web-flow · commit 1f5b8eb366ff · 2024-09-18T10:56:27.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,13 @@
 # @neo4j/neo4j-graphrag-python
 
 ## Next
+### Added
+- Add `template` validation in `PromptTemplate` class upon construction.
+- `custom_prompt` arg is now converted to `Text2CypherTemplate` class within the `Text2CypherRetriever.get_search_results` method.
+- `Text2CypherTemplate` and `RAGTemplate` prompt templates now require `query_text` arg and will error if it is not present. Previous `query_text` aliases may be used, but will warn of deprecation.
+- Fix bug in `Text2CypherRetriever` using `custom_prompt` arg where the `search` method would not inject the `query_text` content.
+- Add feature to include kwargs in `Text2CypherRetriever.search()` that will be injected into a custom prompt, if provided.
+- Add validation to `custom_prompt` parameter of `Text2CypherRetriever` to ensure that `query_text` placeholder exists in prompt.
 
 ## 0.6.3
 ### Changed
@@ -13,10 +20,6 @@
 - Updated documentation to include OpenAI and Vertex AI embeddings classes.
 - Added google-cloud-aiplatform as an optional dependency for Vertex AI embeddings.
 
-### Fixed
-- Make `pygraphviz` an optional dependency - it is now only required when calling `pipeline.draw`.
-
-
 ## 0.6.2
 
 ### Fixed
@@ -125,4 +128,4 @@
 ### Fixed
 
 -   Updated documentation to include new custom exceptions.
--   Improved the use of Pydantic for input data validation for retriever objects.
+-   Improved the use of Pydantic for input data validation for retriever objects.
diff --git a/docs/Makefile b/docs/Makefile
@@ -191,4 +191,4 @@ xml:
 pseudoxml:
 	$(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml
 	@echo
-	@echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
+	@echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml."
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -383,4 +383,4 @@ PipelineStatusUpdateError
 =========================
 
 .. autoclass:: neo4j_graphrag.experimental.pipeline.exceptions.PipelineStatusUpdateError
-   :show-inheritance:
+   :show-inheritance:
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -302,4 +302,4 @@ Indices and tables
 
 * :ref:`genindex`
 * :ref:`modindex`
-* :ref:`search`
+* :ref:`search`
diff --git a/docs/source/types.rst b/docs/source/types.rst
@@ -81,4 +81,4 @@ SchemaRelation
 SchemaConfig
 ============
 
-.. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaConfig
+.. autoclass:: neo4j_graphrag.experimental.components.schema.SchemaConfig
diff --git a/docs/source/user_guide_kg_builder.rst b/docs/source/user_guide_kg_builder.rst
@@ -409,4 +409,4 @@ It is possible to create a custom writer using the `KGWriter` interface:
     The `validate_call` decorator is required when the input parameter contain a `pydantic` model.
 
 
-See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.
+See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.
diff --git a/docs/source/user_guide_pipeline.rst b/docs/source/user_guide_pipeline.rst
@@ -130,4 +130,4 @@ can be added to the canvas by setting `hide_unused_outputs` to `False`:
 Here is an example of final result:
 
 .. image:: images/pipeline_full.png
-  :alt: Pipeline visualisation
+  :alt: Pipeline visualisation
diff --git a/docs/source/user_guide_rag.rst b/docs/source/user_guide_rag.rst
@@ -772,4 +772,4 @@ Drop a Vector Index
 
     # Connect to Neo4j database
     driver = GraphDatabase.driver(URI, auth=AUTH)
-    drop_index_if_exists(driver, INDEX_NAME)
+    drop_index_if_exists(driver, INDEX_NAME)
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -108,4 +108,4 @@ module = [
     "sentence_transformers.*",
     "conftest",
 ]
-ignore_missing_imports = true
+ignore_missing_imports = true
diff --git a/src/neo4j_graphrag/exceptions.py b/src/neo4j_graphrag/exceptions.py
@@ -116,3 +116,7 @@ class PdfLoaderError(Neo4jGraphRagError):
     """Custom exception for errors in PDF loader."""
 
     pass
+
+
+class PromptMissingPlaceholderError(Neo4jGraphRagError):
+    """Exception raised when a prompt is missing an expected placeholder."""
diff --git a/src/neo4j_graphrag/generation/prompts.py b/src/neo4j_graphrag/generation/prompts.py
@@ -14,9 +14,13 @@
 #  limitations under the License.
 from __future__ import annotations
 
+import warnings
 from typing import Any, Optional
 
-from neo4j_graphrag.exceptions import PromptMissingInputError
+from neo4j_graphrag.exceptions import (
+    PromptMissingInputError,
+    PromptMissingPlaceholderError,
+)
 
 
 class PromptTemplate:
@@ -29,7 +33,7 @@ class PromptTemplate:
     """
 
     DEFAULT_TEMPLATE: str = ""
-    EXPECTED_INPUTS: list[str] = []
+    EXPECTED_INPUTS: list[str] = list()
 
     def __init__(
         self,
@@ -39,6 +43,12 @@ def __init__(
         self.template = template or self.DEFAULT_TEMPLATE
         self.expected_inputs = expected_inputs or self.EXPECTED_INPUTS
 
+        for e in self.expected_inputs:
+            if f"{{{e}}}" not in self.template:
+                raise PromptMissingPlaceholderError(
+                    f"`template` is missing placeholder {e}"
+                )
+
     def _format(self, **kwargs: Any) -> str:
         for e in self.EXPECTED_INPUTS:
             if e not in kwargs:
@@ -108,17 +118,41 @@ class Text2CypherTemplate(PromptTemplate):
 {examples}
 
 Input:
-{query}
+{query_text}
 
 Do not use any properties or relationships not included in the schema.
 Do not include triple backticks ``` or any additional text except the generated Cypher statement in your response.
 
 Cypher query:
 """
-    EXPECTED_INPUTS = ["schema", "query", "examples"]
+    EXPECTED_INPUTS = ["query_text"]
 
-    def format(self, query: str, schema: str, examples: str) -> str:
-        return super().format(query=query, schema=schema, examples=examples)
+    def format(
+        self,
+        schema: Optional[str] = None,
+        examples: Optional[str] = None,
+        query_text: str = "",
+        query: Optional[str] = None,
+        **kwargs: Any,
+    ) -> str:
+        if query is not None:
+            if query_text:
+                warnings.warn(
+                    "Both 'query' and 'query_text' are provided, 'query_text' will be used.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+            elif isinstance(query, str):
+                warnings.warn(
+                    "'query' is deprecated and will be removed in a future version, please use 'query_text' instead.",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                query_text = query
+
+        return super().format(
+            query_text=query_text, schema=schema, examples=examples, **kwargs
+        )
 
 
 class ERExtractionTemplate(PromptTemplate):
@@ -151,5 +185,10 @@ class ERExtractionTemplate(PromptTemplate):
 """
     EXPECTED_INPUTS = ["text"]
 
-    def format(self, text: str, schema: dict[str, Any], examples: str) -> str:
+    def format(
+        self,
+        schema: dict[str, Any],
+        examples: str,
+        text: str = "",
+    ) -> str:
         return super().format(text=text, schema=schema, examples=examples)
diff --git a/src/neo4j_graphrag/retrievers/text2cypher.py b/src/neo4j_graphrag/retrievers/text2cypher.py
@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, Optional
 
 import neo4j
 from neo4j.exceptions import CypherSyntaxError, DriverError, Neo4jError
@@ -103,6 +103,8 @@ def __init__(
                     if validated_data.neo4j_schema_model
                     else get_schema(validated_data.driver_model.driver)
                 )
+            else:
+                self.neo4j_schema = ""
 
         except (Neo4jError, DriverError) as e:
             error_message = getattr(e, "message", str(e))
@@ -111,14 +113,14 @@ def __init__(
             ) from e
 
     def get_search_results(
-        self,
-        query_text: str,
+        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. Example: {'schema': 'this is the graph schema'}
 
         Raises:
             SearchValidationError: If validation of the input arguments fail.
@@ -132,15 +134,27 @@ def get_search_results(
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e
 
-        if not self.custom_prompt:
-            prompt_template = Text2CypherTemplate()
-            prompt = prompt_template.format(
-                schema=self.neo4j_schema,
-                examples="\n".join(self.examples) if self.examples else "",
-                query=validated_data.query_text,
+        prompt_template = Text2CypherTemplate(template=self.custom_prompt)
+
+        if prompt_params is not None:
+            # parse the schema and examples inputs
+            examples_to_use = prompt_params.get("examples") or (
+                "\n".join(self.examples) if self.examples else ""
             )
+            schema_to_use = prompt_params.get("schema") or self.neo4j_schema
+            prompt_params.pop("examples", None)
+            prompt_params.pop("schema", None)
         else:
-            prompt = self.custom_prompt
+            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)
 
diff --git a/tests/unit/retrievers/test_text2cypher.py b/tests/unit/retrievers/test_text2cypher.py
diff --git a/tests/unit/test_prompt_template.py b/tests/unit/test_prompt_template.py

Original file line number	Diff line number	Diff line change
@@ -409,4 +409,4 @@ It is possible to create a custom writer using the `KGWriter` interface:
`409`	`409`	The `validate_call` decorator is required when the input parameter contain a `pydantic` model.
`410`	`410`
`411`	`411`
`412`		-See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.
	`412`	+See :ref:`kgwritermodel` and :ref:`kgwriter` in API reference.
Original file line number	Diff line number	Diff line change
`@@ -108,4 +108,4 @@ module = [`
`108`	`108`	`"sentence_transformers.*",`
`109`	`109`	`"conftest",`
`110`	`110`	`]`
`111`		`-ignore_missing_imports = true`
	`111`	`+ignore_missing_imports = true`