Skip to content

Commit

Permalink
Add support for parameters in neo4j retrieval query (#18310)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
tomasonjo authored Feb 29, 2024
1 parent 15d1b73 commit 5999c4a
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 6 deletions.
20 changes: 14 additions & 6 deletions libs/community/langchain_community/vectorstores/neo4j_vector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"}),
]

0 comments on commit 5999c4a

Please sign in to comment.