mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-13 21:47:12 +00:00
Add KuzuQAChain (#6454)
This PR adds `KuzuGraph` and `KuzuQAChain` for interacting with [Kùzu database](https://github.com/kuzudb/kuzu). Kùzu is an in-process property graph database management system (GDBMS) built for query speed and scalability. The `KuzuGraph` and `KuzuQAChain` provide the same functionality as the existing integration with NebulaGraph and Neo4j and enables query generation and question answering over Kùzu database. A notebook example and a simple test case have also been added. --------- Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
This commit is contained in:
56
tests/integration_tests/test_kuzu.py
Normal file
56
tests/integration_tests/test_kuzu.py
Normal file
@@ -0,0 +1,56 @@
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from langchain.graphs import KuzuGraph
|
||||
|
||||
EXPECTED_SCHEMA = """
|
||||
Node properties: [{'properties': [('name', 'STRING')], 'label': 'Movie'}, {'properties': [('name', 'STRING'), ('birthDate', 'STRING')], 'label': 'Person'}]
|
||||
Relationships properties: [{'properties': [], 'label': 'ActedIn'}]
|
||||
Relationships: ['(:Person)-[:ActedIn]->(:Movie)']
|
||||
""" # noqa: E501
|
||||
|
||||
|
||||
class TestKuzu(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
try:
|
||||
import kuzu
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Cannot import Python package kuzu. Please install it by running "
|
||||
"`pip install kuzu`."
|
||||
) from e
|
||||
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.kuzu_database = kuzu.Database(self.tmpdir)
|
||||
self.conn = kuzu.Connection(self.kuzu_database)
|
||||
self.conn.execute("CREATE NODE TABLE Movie (name STRING, PRIMARY KEY(name))")
|
||||
self.conn.execute("CREATE (:Movie {name: 'The Godfather'})")
|
||||
self.conn.execute("CREATE (:Movie {name: 'The Godfather: Part II'})")
|
||||
self.conn.execute(
|
||||
"CREATE (:Movie {name: 'The Godfather Coda: The Death of Michael "
|
||||
"Corleone'})"
|
||||
)
|
||||
self.kuzu_graph = KuzuGraph(self.kuzu_database)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
shutil.rmtree(self.tmpdir, ignore_errors=True)
|
||||
|
||||
def test_query(self) -> None:
|
||||
result = self.kuzu_graph.query("MATCH (n:Movie) RETURN n.name ORDER BY n.name")
|
||||
excepted_result = [
|
||||
{"n.name": "The Godfather"},
|
||||
{"n.name": "The Godfather Coda: The Death of Michael Corleone"},
|
||||
{"n.name": "The Godfather: Part II"},
|
||||
]
|
||||
self.assertEqual(result, excepted_result)
|
||||
|
||||
def test_refresh_schema(self) -> None:
|
||||
self.conn.execute(
|
||||
"CREATE NODE TABLE Person (name STRING, birthDate STRING, PRIMARY "
|
||||
"KEY(name))"
|
||||
)
|
||||
self.conn.execute("CREATE REL TABLE ActedIn (FROM Person TO Movie)")
|
||||
self.kuzu_graph.refresh_schema()
|
||||
schema = self.kuzu_graph.get_schema
|
||||
self.assertEqual(schema, EXPECTED_SCHEMA)
|
Reference in New Issue
Block a user