mirror of
https://github.com/hwchase17/langchain.git
synced 2025-06-20 05:43:55 +00:00
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:
parent
a84a3add25
commit
bccc9241ea
@ -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 = (
|
||||||
|
@ -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 "
|
||||||
|
Loading…
Reference in New Issue
Block a user