mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-04 04:07:54 +00:00
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.
This commit is contained in:
parent
15d1b73a00
commit
5999c4a240
@ -549,6 +549,7 @@ class Neo4jVector(VectorStore):
|
||||
self,
|
||||
query: str,
|
||||
k: int = 4,
|
||||
params: Dict[str, Any] = {},
|
||||
**kwargs: Any,
|
||||
) -> List[Document]:
|
||||
"""Run similarity search with Neo4jVector.
|
||||
@ -562,13 +563,15 @@ class Neo4jVector(VectorStore):
|
||||
"""
|
||||
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 @@ class Neo4jVector(VectorStore):
|
||||
"""
|
||||
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 @@ class Neo4jVector(VectorStore):
|
||||
"embedding": embedding,
|
||||
"keyword_index": self.keyword_index_name,
|
||||
"query": remove_lucene_chars(kwargs["query"]),
|
||||
**params,
|
||||
}
|
||||
|
||||
results = self.query(read_query, params=parameters)
|
||||
|
@ -721,3 +721,23 @@ def test_index_fetching() -> None:
|
||||
|
||||
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"}),
|
||||
]
|
||||
|
Loading…
Reference in New Issue
Block a user