Move few-shot examples in Text2CypherRetriever to constructor (#62)

* Moved the few-shot examples in the Text2CypherRetriever to the constructor

* Updated CHANGELOG

Author: alexthomas93

CHANGELOG.md

@@ -12,6 +12,7 @@
 ### Fixed
 
 -   Removed Pinecone and Weaviate retrievers from **init**.py to prevent ImportError when optional dependencies are not installed.
+-   Moved few-shot examples in `Text2CypherRetriever` to the constructor for better initialization and usage. Updated unit tests and example script accordingly.
 
 ## 0.2.0a5
 

examples/text2cypher_search.py

@@ -28,14 +28,19 @@
 (:Person)-[:REVIEWED]->(:Movie)
 """
 
-# Initialize the retriever
-retriever = Text2CypherRetriever(driver=driver, llm=llm, neo4j_schema=neo4j_schema)  # type: ignore
-
 # (Optional) Provide user input/query pairs for the LLM to use as examples
 examples = [
     "USER INPUT: 'Which actors starred in the Matrix?' QUERY: MATCH (p:Person)-[:ACTED_IN]->(m:Movie) WHERE m.title = 'The Matrix' RETURN p.name"
 ]
 
+# Initialize the retriever
+retriever = Text2CypherRetriever(
+    driver=driver,
+    llm=llm,  # type: ignore
+    neo4j_schema=neo4j_schema,
+    examples=examples,
+)
+
 # Generate a Cypher query using the LLM, send it to the Neo4j database, and return the results
 query_text = "Which movies did Hugo Weaving star in?"
-print(retriever.search(query_text=query_text, examples=examples))
+print(retriever.search(query_text=query_text))

src/neo4j_genai/retrievers/text2cypher.py

@@ -21,9 +21,9 @@
 
 from neo4j_genai.exceptions import (
     RetrieverInitializationError,
+    SchemaFetchError,
     SearchValidationError,
     Text2CypherRetrievalError,
-    SchemaFetchError,
 )
 from neo4j_genai.llm import LLM
 from neo4j_genai.prompts import TEXT2CYPHER_PROMPT
@@ -33,9 +33,9 @@
     LLMModel,
     Neo4jDriverModel,
     Neo4jSchemaModel,
+    RawSearchResult,
     Text2CypherRetrieverModel,
     Text2CypherSearchModel,
-    RawSearchResult,
 )
 
 logger = logging.getLogger(__name__)
@@ -51,13 +51,18 @@ class Text2CypherRetriever(Retriever):
         driver (neo4j.driver): The Neo4j Python driver.
         llm (neo4j_genai.llm.LLM): 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.
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
     """
 
     def __init__(
-        self, driver: neo4j.Driver, llm: LLM, neo4j_schema: Optional[str] = None
+        self,
+        driver: neo4j.Driver,
+        llm: LLM,
+        neo4j_schema: Optional[str] = None,
+        examples: Optional[list[str]] = None,
     ) -> None:
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -69,12 +74,14 @@ def __init__(
                 driver_model=driver_model,
                 llm_model=llm_model,
                 neo4j_schema_model=neo4j_schema_model,
+                examples=examples,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors())
 
         super().__init__(validated_data.driver_model.driver)
         self.llm = validated_data.llm_model.llm
+        self.examples = validated_data.examples
         try:
             self.neo4j_schema = (
                 validated_data.neo4j_schema_model.neo4j_schema
@@ -88,14 +95,14 @@ def __init__(
             )
 
     def _get_search_results(
-        self, query_text: str, examples: Optional[list[str]] = None
+        self,
+        query_text: str,
     ) -> 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.
-            examples (Optional[list[str], optional): Optional user input/query pairs for the LLM to use as examples.
 
         Raises:
             SearchValidationError: If validation of the input arguments fail.
@@ -105,17 +112,13 @@ def _get_search_results(
             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, examples=examples
-            )
+            validated_data = Text2CypherSearchModel(query_text=query_text)
         except ValidationError as e:
             raise SearchValidationError(e.errors())
 
         prompt = TEXT2CYPHER_PROMPT.format(
             schema=self.neo4j_schema,
-            examples="\n".join(validated_data.examples)
-            if validated_data.examples
-            else "",
+            examples="\n".join(self.examples) if self.examples else "",
             input=validated_data.query_text,
         )
         logger.debug("Text2CypherRetriever prompt: %s", prompt)

src/neo4j_genai/types.py

@@ -148,7 +148,6 @@ class HybridCypherSearchModel(HybridSearchModel):
 
 class Text2CypherSearchModel(BaseModel):
     query_text: str
-    examples: Optional[list[str]] = None
 
 
 class SearchType(str, Enum):
@@ -231,3 +230,4 @@ class Text2CypherRetrieverModel(BaseModel):
     driver_model: Neo4jDriverModel
     llm_model: LLMModel
     neo4j_schema_model: Optional[Neo4jSchemaModel] = None
+    examples: Optional[list[str]] = None

tests/unit/retrievers/test_text2cypher.py

@@ -12,14 +12,14 @@
 #  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 unittest.mock import patch, MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from neo4j.exceptions import CypherSyntaxError, Neo4jError
 from neo4j_genai import Text2CypherRetriever
 from neo4j_genai.exceptions import (
-    SearchValidationError,
     RetrieverInitializationError,
+    SearchValidationError,
     Text2CypherRetrievalError,
 )
 from neo4j_genai.prompts import TEXT2CYPHER_PROMPT
@@ -85,14 +85,16 @@ def test_t2c_retriever_invalid_search_query(
 def test_t2c_retriever_invalid_search_examples(
     _verify_version_mock: MagicMock, driver: MagicMock, llm: MagicMock
 ) -> None:
-    with pytest.raises(SearchValidationError) as exc_info:
-        retriever = Text2CypherRetriever(
-            driver=driver, llm=llm, neo4j_schema="dummy-text"
+    with pytest.raises(RetrieverInitializationError) as exc_info:
+        Text2CypherRetriever(
+            driver=driver,
+            llm=llm,
+            neo4j_schema="dummy-text",
+            examples=42,  # type: ignore
         )
-        retriever.search(query_text="dummy-text", examples=42)
 
     assert "examples" in str(exc_info.value)
-    assert "Input should be a valid list" in str(exc_info.value)
+    assert "Initialization failed" in str(exc_info.value)
 
 
 @patch("neo4j_genai.Text2CypherRetriever._verify_version")
@@ -106,7 +108,9 @@ def test_t2c_retriever_happy_path(
     query_text = "may thy knife chip and shatter"
     neo4j_schema = "dummy-schema"
     examples = ["example-1", "example-2"]
-    retriever = Text2CypherRetriever(driver=driver, llm=llm, neo4j_schema=neo4j_schema)
+    retriever = Text2CypherRetriever(
+        driver=driver, llm=llm, neo4j_schema=neo4j_schema, examples=examples
+    )
     retriever.llm.invoke.return_value = t2c_query
     retriever.driver.execute_query.return_value = (  # type: ignore
         [neo4j_record],
@@ -118,7 +122,7 @@ def test_t2c_retriever_happy_path(
         examples="\n".join(examples),
         input=query_text,
     )
-    retriever.search(query_text=query_text, examples=examples)
+    retriever.search(query_text=query_text)
     retriever.llm.invoke.assert_called_once_with(prompt)
     retriever.driver.execute_query.assert_called_once_with(query_=t2c_query)  # type: ignore
 
@@ -130,10 +134,12 @@ def test_t2c_retriever_cypher_error(
     t2c_query = "this is not a cypher query"
     neo4j_schema = "dummy-schema"
     examples = ["example-1", "example-2"]
-    retriever = Text2CypherRetriever(driver=driver, llm=llm, neo4j_schema=neo4j_schema)
+    retriever = Text2CypherRetriever(
+        driver=driver, llm=llm, neo4j_schema=neo4j_schema, examples=examples
+    )
     retriever.llm.invoke.return_value = t2c_query
     query_text = "may thy knife chip and shatter"
     driver.execute_query.side_effect = CypherSyntaxError
     with pytest.raises(Text2CypherRetrievalError) as e:
-        retriever.search(query_text=query_text, examples=examples)
+        retriever.search(query_text=query_text)
     assert "Failed to get search result" in str(e)