community[patch]: Add neo4j timeout and value sanitization option (#16138)

The timeout function comes in handy when you want to kill longrunning
queries.
The value sanitization removes all lists that are larger than 128
elements. The idea here is to remove embedding properties from results.
This commit is contained in:
Tomaz Bratanic
2024-01-17 22:22:19 +01:00
committed by GitHub
parent 27ed2673da
commit 1e80113ac9
3 changed files with 104 additions and 3 deletions

View File

@@ -69,3 +69,22 @@ def test_cypher_return_correct_schema() -> None:
sorted(relationships, key=lambda x: x["output"]["end"])
== expected_relationships
)
def test_neo4j_timeout() -> None:
"""Test that neo4j uses the timeout correctly."""
url = os.environ.get("NEO4J_URI")
username = os.environ.get("NEO4J_USERNAME")
password = os.environ.get("NEO4J_PASSWORD")
assert url is not None
assert username is not None
assert password is not None
graph = Neo4jGraph(url=url, username=username, password=password, timeout=0.1)
try:
graph.query("UNWIND range(0,100000,1) AS i MERGE (:Foo {id:i})")
except Exception as e:
assert (
e.code
== "Neo.ClientError.Transaction.TransactionTimedOutClientConfiguration"
)

View File

@@ -0,0 +1,32 @@
from langchain_community.graphs.neo4j_graph import value_sanitize
def test_value_sanitize_with_small_list():
small_list = list(range(15)) # list size > LIST_LIMIT
input_dict = {"key1": "value1", "small_list": small_list}
expected_output = {"key1": "value1", "small_list": small_list}
assert value_sanitize(input_dict) == expected_output
def test_value_sanitize_with_oversized_list():
oversized_list = list(range(150)) # list size > LIST_LIMIT
input_dict = {"key1": "value1", "oversized_list": oversized_list}
expected_output = {
"key1": "value1"
# oversized_list should not be included
}
assert value_sanitize(input_dict) == expected_output
def test_value_sanitize_with_nested_oversized_list():
oversized_list = list(range(150)) # list size > LIST_LIMIT
input_dict = {"key1": "value1", "oversized_list": {"key": oversized_list}}
expected_output = {"key1": "value1", "oversized_list": {}}
assert value_sanitize(input_dict) == expected_output
def test_value_sanitize_with_dict_in_list():
oversized_list = list(range(150)) # list size > LIST_LIMIT
input_dict = {"key1": "value1", "oversized_list": [1, 2, {"key": oversized_list}]}
expected_output = {"key1": "value1", "oversized_list": [1, 2, {}]}
assert value_sanitize(input_dict) == expected_output