From 5999c4a24093e23f3d647bf861259578cc19bc3d Mon Sep 17 00:00:00 2001
From: Tomaz Bratanic <bratanic.tomaz@gmail.com>
Date: Thu, 29 Feb 2024 22:00:54 +0100
Subject: [PATCH] Add support for parameters in neo4j retrieval query (#18310)

Sometimes, you want to use various parameters in the retrieval query of
Neo4j Vector to personalize/customize results. Before, when there were
only predefined chains, it didn't really make sense. Now that it's all
about custom chains and LCEL, it is worth adding since users can inject
any params they wish at query time. Isn't prone to SQL injection-type
attacks since we use parameters and not concatenating strings.
---
 .../vectorstores/neo4j_vector.py              | 20 +++++++++++++------
 .../vectorstores/test_neo4jvector.py          | 20 +++++++++++++++++++
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/libs/community/langchain_community/vectorstores/neo4j_vector.py b/libs/community/langchain_community/vectorstores/neo4j_vector.py
index 5e9057c533dae..2ce1366ba6085 100644
--- a/libs/community/langchain_community/vectorstores/neo4j_vector.py
+++ b/libs/community/langchain_community/vectorstores/neo4j_vector.py
@@ -549,6 +549,7 @@ def similarity_search(
         self,
         query: str,
         k: int = 4,
+        params: Dict[str, Any] = {},
         **kwargs: Any,
     ) -> List[Document]:
         """Run similarity search with Neo4jVector.
@@ -562,13 +563,15 @@ def similarity_search(
         """
         embedding = self.embedding.embed_query(text=query)
         return self.similarity_search_by_vector(
-            embedding=embedding,
-            k=k,
-            query=query,
+            embedding=embedding, k=k, query=query, params=params, **kwargs
         )
 
     def similarity_search_with_score(
-        self, query: str, k: int = 4
+        self,
+        query: str,
+        k: int = 4,
+        params: Dict[str, Any] = {},
+        **kwargs: Any,
     ) -> List[Tuple[Document, float]]:
         """Return docs most similar to query.
 
@@ -581,12 +584,16 @@ def similarity_search_with_score(
         """
         embedding = self.embedding.embed_query(query)
         docs = self.similarity_search_with_score_by_vector(
-            embedding=embedding, k=k, query=query
+            embedding=embedding, k=k, query=query, params=params, **kwargs
         )
         return docs
 
     def similarity_search_with_score_by_vector(
-        self, embedding: List[float], k: int = 4, **kwargs: Any
+        self,
+        embedding: List[float],
+        k: int = 4,
+        params: Dict[str, Any] = {},
+        **kwargs: Any,
     ) -> List[Tuple[Document, float]]:
         """
         Perform a similarity search in the Neo4j database using a
@@ -623,6 +630,7 @@ def similarity_search_with_score_by_vector(
             "embedding": embedding,
             "keyword_index": self.keyword_index_name,
             "query": remove_lucene_chars(kwargs["query"]),
+            **params,
         }
 
         results = self.query(read_query, params=parameters)
diff --git a/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py b/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py
index 13aef52a40859..ffa340398f563 100644
--- a/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py
+++ b/libs/community/tests/integration_tests/vectorstores/test_neo4jvector.py
@@ -721,3 +721,23 @@ def fetch_store(index_name: str) -> Neo4jVector:
 
     index_0_store = fetch_store(index_0_str)
     assert index_0_store.index_name == index_0_str
+
+
+def test_retrieval_params() -> None:
+    """Test if we use parameters in retrieval query"""
+    docsearch = Neo4jVector.from_texts(
+        texts=texts,
+        embedding=FakeEmbeddings(),
+        pre_delete_collection=True,
+        retrieval_query="""
+        RETURN $test as text, score, {test: $test1} AS metadata
+        """,
+    )
+
+    output = docsearch.similarity_search(
+        "Foo", k=2, params={"test": "test", "test1": "test1"}
+    )
+    assert output == [
+        Document(page_content="test", metadata={"test": "test1"}),
+        Document(page_content="test", metadata={"test": "test1"}),
+    ]