Skip to content
Snippets Groups Projects
Unverified Commit f3915f3e authored by Egor Kraev's avatar Egor Kraev Committed by GitHub
Browse files

Fix kuzu integration .execute() calls, which require a dict not a list of tuples (#13100)

* Fix kuzu execute calls, which require a dict not a list of tuples

* Linter fix

* Increment version in pyproject.toml

* Increment version in the correct pyproject.toml, roll back the other one
parent 9a917864
No related branches found
No related tags found
No related merge requests found
"""Kùzu graph store index.""" """Kùzu graph store index."""
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from llama_index.core.graph_stores.types import GraphStore from llama_index.core.graph_stores.types import GraphStore
...@@ -51,7 +52,7 @@ class KuzuGraphStore(GraphStore): ...@@ -51,7 +52,7 @@ class KuzuGraphStore(GraphStore):
prepared_statement = self.connection.prepare( prepared_statement = self.connection.prepare(
query % (self.node_table_name, self.rel_table_name, self.node_table_name) query % (self.node_table_name, self.rel_table_name, self.node_table_name)
) )
query_result = self.connection.execute(prepared_statement, [("subj", subj)]) query_result = self.connection.execute(prepared_statement, {"subj": subj})
retval = [] retval = []
while query_result.has_next(): while query_result.has_next():
row = query_result.get_next() row = query_result.get_next()
...@@ -82,7 +83,7 @@ class KuzuGraphStore(GraphStore): ...@@ -82,7 +83,7 @@ class KuzuGraphStore(GraphStore):
query = f"{match_clause} {where_clause} {return_clause}" query = f"{match_clause} {where_clause} {return_clause}"
prepared_statement = self.connection.prepare(query) prepared_statement = self.connection.prepare(query)
if subjs is not None: if subjs is not None:
query_result = self.connection.execute(prepared_statement, params) query_result = self.connection.execute(prepared_statement, dict(params))
else: else:
query_result = self.connection.execute(prepared_statement) query_result = self.connection.execute(prepared_statement)
retval: Dict[str, List[List[str]]] = {} retval: Dict[str, List[List[str]]] = {}
...@@ -115,14 +116,14 @@ class KuzuGraphStore(GraphStore): ...@@ -115,14 +116,14 @@ class KuzuGraphStore(GraphStore):
def check_entity_exists(connection: Any, entity: str) -> bool: def check_entity_exists(connection: Any, entity: str) -> bool:
is_exists_result = connection.execute( is_exists_result = connection.execute(
"MATCH (n:%s) WHERE n.ID = $entity RETURN n.ID" % self.node_table_name, "MATCH (n:%s) WHERE n.ID = $entity RETURN n.ID" % self.node_table_name,
[("entity", entity)], {"entity": entity},
) )
return is_exists_result.has_next() return is_exists_result.has_next()
def create_entity(connection: Any, entity: str) -> None: def create_entity(connection: Any, entity: str) -> None:
connection.execute( connection.execute(
"CREATE (n:%s {ID: $entity})" % self.node_table_name, "CREATE (n:%s {ID: $entity})" % self.node_table_name,
[("entity", entity)], {"entity": entity},
) )
def check_rel_exists(connection: Any, subj: str, obj: str, rel: str) -> bool: def check_rel_exists(connection: Any, subj: str, obj: str, rel: str) -> bool:
...@@ -133,7 +134,7 @@ class KuzuGraphStore(GraphStore): ...@@ -133,7 +134,7 @@ class KuzuGraphStore(GraphStore):
).format( ).format(
self.node_table_name, self.rel_table_name, self.node_table_name self.node_table_name, self.rel_table_name, self.node_table_name
), ),
[("subj", subj), ("obj", obj), ("pred", rel)], {"subj": subj, "obj": obj, "pred": rel},
) )
return is_exists_result.has_next() return is_exists_result.has_next()
...@@ -145,7 +146,7 @@ class KuzuGraphStore(GraphStore): ...@@ -145,7 +146,7 @@ class KuzuGraphStore(GraphStore):
).format( ).format(
self.node_table_name, self.node_table_name, self.rel_table_name self.node_table_name, self.node_table_name, self.rel_table_name
), ),
[("subj", subj), ("obj", obj), ("pred", rel)], {"subj": subj, "obj": obj, "pred": rel},
) )
is_subj_exists = check_entity_exists(self.connection, subj) is_subj_exists = check_entity_exists(self.connection, subj)
...@@ -174,13 +175,13 @@ class KuzuGraphStore(GraphStore): ...@@ -174,13 +175,13 @@ class KuzuGraphStore(GraphStore):
).format( ).format(
self.node_table_name, self.rel_table_name, self.node_table_name self.node_table_name, self.rel_table_name, self.node_table_name
), ),
[("subj", subj), ("obj", obj), ("pred", rel)], {"subj": subj, "obj": obj, "pred": rel},
) )
def delete_entity(connection: Any, entity: str) -> None: def delete_entity(connection: Any, entity: str) -> None:
connection.execute( connection.execute(
"MATCH (n:%s) WHERE n.ID = $entity DELETE n" % self.node_table_name, "MATCH (n:%s) WHERE n.ID = $entity DELETE n" % self.node_table_name,
[("entity", entity)], {"entity": entity},
) )
def check_edges(connection: Any, entity: str) -> bool: def check_edges(connection: Any, entity: str) -> bool:
...@@ -188,7 +189,7 @@ class KuzuGraphStore(GraphStore): ...@@ -188,7 +189,7 @@ class KuzuGraphStore(GraphStore):
"MATCH (n1:{})-[r:{}]-(n2:{}) WHERE n2.ID = $entity RETURN r.predicate".format( "MATCH (n1:{})-[r:{}]-(n2:{}) WHERE n2.ID = $entity RETURN r.predicate".format(
self.node_table_name, self.rel_table_name, self.node_table_name self.node_table_name, self.rel_table_name, self.node_table_name
), ),
[("entity", entity)], {"entity": entity},
) )
return is_exists_result.has_next() return is_exists_result.has_next()
......
...@@ -27,7 +27,7 @@ exclude = ["**/BUILD"] ...@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
license = "MIT" license = "MIT"
name = "llama-index-graph-stores-kuzu" name = "llama-index-graph-stores-kuzu"
readme = "README.md" readme = "README.md"
version = "0.1.2" version = "0.1.3"
[tool.poetry.dependencies] [tool.poetry.dependencies]
python = ">=3.8.1,<4.0" python = ">=3.8.1,<4.0"
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment