community[patch]: Resolve KuzuQAChain API Changes (#16885)

- **Description:** Updates to the Kuzu API had broken this
functionality. These updates resolve those issues and add a new test to
demonstrate the updates.
- **Issue:** #11874
- **Dependencies:** No new dependencies
- **Twitter handle:** @amirk08


Test results:
```
tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_query_no_params PASSED                                   [ 33%]
tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_query_params PASSED                                      [ 66%]
tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_refresh_schema PASSED                                    [100%]

=================================================== slowest 5 durations =================================================== 
0.53s call     tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_refresh_schema
0.34s call     tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_query_no_params
0.28s call     tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_query_params
0.03s teardown tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_refresh_schema
0.02s teardown tests/integration_tests/graphs/test_kuzu.py::TestKuzu::test_query_params
==================================================== 3 passed in 1.27s ==================================================== 
```
This commit is contained in:
Amir Karbasi 2024-02-15 13:18:37 -05:00 committed by GitHub
parent a84a3add25
commit bccc9241ea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 21 deletions

View File

@ -36,10 +36,7 @@ class KuzuGraph:
def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]: def query(self, query: str, params: dict = {}) -> List[Dict[str, Any]]:
"""Query Kùzu database""" """Query Kùzu database"""
params_list = [] result = self.conn.execute(query, params)
for param_name in params:
params_list.append([param_name, params[param_name]])
result = self.conn.execute(query, params_list)
column_names = result.get_column_names() column_names = result.get_column_names()
return_list = [] return_list = []
while result.has_next(): while result.has_next():
@ -79,20 +76,16 @@ class KuzuGraph:
rel_properties = [] rel_properties = []
for table in rel_tables: for table in rel_tables:
current_table_schema = {"properties": [], "label": table["name"]} table_name = table["name"]
properties_text = self.conn._connection.get_rel_property_names( current_table_schema = {"properties": [], "label": table_name}
table["name"] query_result = self.conn.execute(
).split("\n") f"CALL table_info('{table_name}') RETURN *;"
for i, line in enumerate(properties_text): )
# The first 3 lines defines src, dst and name, so we skip them while query_result.has_next():
if i < 3: row = query_result.get_next()
continue prop_name = row[1]
if not line: prop_type = row[2]
continue current_table_schema["properties"].append((prop_name, prop_type))
property_name, property_type = line.strip().split(" ")
current_table_schema["properties"].append(
(property_name, property_type)
)
rel_properties.append(current_table_schema) rel_properties.append(current_table_schema)
self.schema = ( self.schema = (

View File

@ -4,8 +4,7 @@ import unittest
from langchain_community.graphs import KuzuGraph from langchain_community.graphs import KuzuGraph
EXPECTED_SCHEMA = """ EXPECTED_SCHEMA = """Node properties: [{'properties': [('name', 'STRING')], 'label': 'Movie'}, {'properties': [('name', 'STRING'), ('birthDate', 'STRING')], 'label': 'Person'}]
Node properties: [{'properties': [('name', 'STRING')], 'label': 'Movie'}, {'properties': [('name', 'STRING'), ('birthDate', 'STRING')], 'label': 'Person'}]
Relationships properties: [{'properties': [], 'label': 'ActedIn'}] Relationships properties: [{'properties': [], 'label': 'ActedIn'}]
Relationships: ['(:Person)-[:ActedIn]->(:Movie)'] Relationships: ['(:Person)-[:ActedIn]->(:Movie)']
""" # noqa: E501 """ # noqa: E501
@ -36,7 +35,7 @@ class TestKuzu(unittest.TestCase):
def tearDown(self) -> None: def tearDown(self) -> None:
shutil.rmtree(self.tmpdir, ignore_errors=True) shutil.rmtree(self.tmpdir, ignore_errors=True)
def test_query(self) -> None: def test_query_no_params(self) -> None:
result = self.kuzu_graph.query("MATCH (n:Movie) RETURN n.name ORDER BY n.name") result = self.kuzu_graph.query("MATCH (n:Movie) RETURN n.name ORDER BY n.name")
excepted_result = [ excepted_result = [
{"n.name": "The Godfather"}, {"n.name": "The Godfather"},
@ -45,6 +44,16 @@ class TestKuzu(unittest.TestCase):
] ]
self.assertEqual(result, excepted_result) self.assertEqual(result, excepted_result)
def test_query_params(self) -> None:
result = self.kuzu_graph.query(
query="MATCH (n:Movie) WHERE n.name = $name RETURN n.name",
params={"name": "The Godfather"},
)
excepted_result = [
{"n.name": "The Godfather"},
]
self.assertEqual(result, excepted_result)
def test_refresh_schema(self) -> None: def test_refresh_schema(self) -> None:
self.conn.execute( self.conn.execute(
"CREATE NODE TABLE Person (name STRING, birthDate STRING, PRIMARY " "CREATE NODE TABLE Person (name STRING, birthDate STRING, PRIMARY "