graphiti-core 0.19.0rc3__py3-none-any.whl → 0.20.0__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -27,10 +27,13 @@ logger = logging.getLogger(__name__)
 class GraphProvider(Enum):
     NEO4J = 'neo4j'
     FALKORDB = 'falkordb'
+    KUZU = 'kuzu'
     NEPTUNE = 'neptune'
 
 
 class GraphDriverSession(ABC):
+    provider: GraphProvider
+
     async def __aenter__(self):
         return self
 

graphiti_core/driver/falkordb_driver.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 import logging
-from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -33,11 +32,14 @@ else:
         ) from None
 
 from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
 
 logger = logging.getLogger(__name__)
 
 
 class FalkorDriverSession(GraphDriverSession):
+    provider = GraphProvider.FALKORDB
+
     def __init__(self, graph: FalkorGraph):
         self.graph = graph
 
@@ -164,16 +166,3 @@ class FalkorDriver(GraphDriver):
         cloned = FalkorDriver(falkor_db=self.client, database=database)
 
         return cloned
-
-
-def convert_datetimes_to_strings(obj):
-    if isinstance(obj, dict):
-        return {k: convert_datetimes_to_strings(v) for k, v in obj.items()}
-    elif isinstance(obj, list):
-        return [convert_datetimes_to_strings(item) for item in obj]
-    elif isinstance(obj, tuple):
-        return tuple(convert_datetimes_to_strings(item) for item in obj)
-    elif isinstance(obj, datetime):
-        return obj.isoformat()
-    else:
-        return obj

graphiti_core/driver/kuzu_driver.py

@@ -0,0 +1,175 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+from typing import Any
+
+import kuzu
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+
+logger = logging.getLogger(__name__)
+
+# Kuzu requires an explicit schema.
+# As Kuzu currently does not support creating full text indexes on edge properties,
+# we work around this by representing (n:Entity)-[:RELATES_TO]->(m:Entity) as
+# (n)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m).
+SCHEMA_QUERIES = """
+    CREATE NODE TABLE IF NOT EXISTS Episodic (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        created_at TIMESTAMP,
+        source STRING,
+        source_description STRING,
+        content STRING,
+        valid_at TIMESTAMP,
+        entity_edges STRING[]
+    );
+    CREATE NODE TABLE IF NOT EXISTS Entity (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        labels STRING[],
+        created_at TIMESTAMP,
+        name_embedding FLOAT[],
+        summary STRING,
+        attributes STRING
+    );
+    CREATE NODE TABLE IF NOT EXISTS Community (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        created_at TIMESTAMP,
+        name_embedding FLOAT[],
+        summary STRING
+    );
+    CREATE NODE TABLE IF NOT EXISTS RelatesToNode_ (
+        uuid STRING PRIMARY KEY,
+        group_id STRING,
+        created_at TIMESTAMP,
+        name STRING,
+        fact STRING,
+        fact_embedding FLOAT[],
+        episodes STRING[],
+        expired_at TIMESTAMP,
+        valid_at TIMESTAMP,
+        invalid_at TIMESTAMP,
+        attributes STRING
+    );
+    CREATE REL TABLE IF NOT EXISTS RELATES_TO(
+        FROM Entity TO RelatesToNode_,
+        FROM RelatesToNode_ TO Entity
+    );
+    CREATE REL TABLE IF NOT EXISTS MENTIONS(
+        FROM Episodic TO Entity,
+        uuid STRING PRIMARY KEY,
+        group_id STRING,
+        created_at TIMESTAMP
+    );
+    CREATE REL TABLE IF NOT EXISTS HAS_MEMBER(
+        FROM Community TO Entity,
+        FROM Community TO Community,
+        uuid STRING,
+        group_id STRING,
+        created_at TIMESTAMP
+    );
+"""
+
+
+class KuzuDriver(GraphDriver):
+    provider: GraphProvider = GraphProvider.KUZU
+
+    def __init__(
+        self,
+        db: str = ':memory:',
+        max_concurrent_queries: int = 1,
+    ):
+        super().__init__()
+        self.db = kuzu.Database(db)
+
+        self.setup_schema()
+
+        self.client = kuzu.AsyncConnection(self.db, max_concurrent_queries=max_concurrent_queries)
+
+    async def execute_query(
+        self, cypher_query_: str, **kwargs: Any
+    ) -> tuple[list[dict[str, Any]] | list[list[dict[str, Any]]], None, None]:
+        params = {k: v for k, v in kwargs.items() if v is not None}
+        # Kuzu does not support these parameters.
+        params.pop('database_', None)
+        params.pop('routing_', None)
+
+        try:
+            results = await self.client.execute(cypher_query_, parameters=params)
+        except Exception as e:
+            params = {k: (v[:5] if isinstance(v, list) else v) for k, v in params.items()}
+            logger.error(f'Error executing Kuzu query: {e}\n{cypher_query_}\n{params}')
+            raise
+
+        if not results:
+            return [], None, None
+
+        if isinstance(results, list):
+            dict_results = [list(result.rows_as_dict()) for result in results]
+        else:
+            dict_results = list(results.rows_as_dict())
+        return dict_results, None, None  # type: ignore
+
+    def session(self, _database: str | None = None) -> GraphDriverSession:
+        return KuzuDriverSession(self)
+
+    async def close(self):
+        # Do not explicity close the connection, instead rely on GC.
+        pass
+
+    def delete_all_indexes(self, database_: str):
+        pass
+
+    def setup_schema(self):
+        conn = kuzu.Connection(self.db)
+        conn.execute(SCHEMA_QUERIES)
+        conn.close()
+
+
+class KuzuDriverSession(GraphDriverSession):
+    provider = GraphProvider.KUZU
+
+    def __init__(self, driver: KuzuDriver):
+        self.driver = driver
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Kuzu, but method must exist.
+        pass
+
+    async def close(self):
+        # Do not close the session here, as we're reusing the driver connection.
+        pass
+
+    async def execute_write(self, func, *args, **kwargs):
+        # Directly await the provided async function with `self` as the transaction/session
+        return await func(self, *args, **kwargs)
+
+    async def run(self, query: str | list, **kwargs: Any) -> Any:
+        if isinstance(query, list):
+            for cypher, params in query:
+                await self.driver.execute_query(cypher, **params)
+        else:
+            await self.driver.execute_query(query, **kwargs)
+        return None

graphiti_core/driver/neptune_driver.py

@@ -271,6 +271,8 @@ class NeptuneDriver(GraphDriver):
 
 
 class NeptuneDriverSession(GraphDriverSession):
+    provider = GraphProvider.NEPTUNE
+
     def __init__(self, driver: NeptuneDriver):  # type: ignore[reportUnknownArgumentType]
         self.driver = driver
 

graphiti_core/edges.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
@@ -30,11 +31,10 @@ from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
 from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.edges.edge_db_queries import (
     COMMUNITY_EDGE_RETURN,
-    ENTITY_EDGE_RETURN,
-    ENTITY_EDGE_RETURN_NEPTUNE,
     EPISODIC_EDGE_RETURN,
     EPISODIC_EDGE_SAVE,
     get_community_edge_save_query,
+    get_entity_edge_return_query,
     get_entity_edge_save_query,
 )
 from graphiti_core.nodes import Node
@@ -53,33 +53,63 @@ class Edge(BaseModel, ABC):
     async def save(self, driver: GraphDriver): ...
 
     async def delete(self, driver: GraphDriver):
-        result = await driver.execute_query(
-            """
-            MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
-            DELETE e
-            """,
-            uuid=self.uuid,
-        )
+        if driver.provider == GraphProvider.KUZU:
+            await driver.execute_query(
+                """
+                MATCH (n)-[e:MENTIONS|HAS_MEMBER {uuid: $uuid}]->(m)
+                DELETE e
+                """,
+                uuid=self.uuid,
+            )
+            await driver.execute_query(
+                """
+                MATCH (e:RelatesToNode_ {uuid: $uuid})
+                DETACH DELETE e
+                """,
+                uuid=self.uuid,
+            )
+        else:
+            await driver.execute_query(
+                """
+                MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
+                DELETE e
+                """,
+                uuid=self.uuid,
+            )
 
         logger.debug(f'Deleted Edge: {self.uuid}')
 
-        return result
-
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
-        result = await driver.execute_query(
-            """
-            MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER]->(m)
-            WHERE e.uuid IN $uuids
-            DELETE e
-            """,
-            uuids=uuids,
-        )
+        if driver.provider == GraphProvider.KUZU:
+            await driver.execute_query(
+                """
+                MATCH (n)-[e:MENTIONS|HAS_MEMBER]->(m)
+                WHERE e.uuid IN $uuids
+                DELETE e
+                """,
+                uuids=uuids,
+            )
+            await driver.execute_query(
+                """
+                MATCH (e:RelatesToNode_)
+                WHERE e.uuid IN $uuids
+                DETACH DELETE e
+                """,
+                uuids=uuids,
+            )
+        else:
+            await driver.execute_query(
+                """
+                MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER]->(m)
+                WHERE e.uuid IN $uuids
+                DELETE e
+                """,
+                uuids=uuids,
+            )
 
         logger.debug(f'Deleted Edges: {uuids}')
 
-        return result
-
     def __hash__(self):
         return hash(self.uuid)
 
@@ -166,7 +196,7 @@ class EpisodicEdge(Edge):
             """
             + EPISODIC_EDGE_RETURN
             + """
-            ORDER BY e.uuid DESC 
+            ORDER BY e.uuid DESC
             """
             + limit_query,
             group_ids=group_ids,
@@ -215,15 +245,21 @@ class EntityEdge(Edge):
         return self.fact_embedding
 
     async def load_fact_embedding(self, driver: GraphDriver):
-        if driver.provider == GraphProvider.NEPTUNE:
-            query: LiteralString = """
+        query = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
+            RETURN e.fact_embedding AS fact_embedding
+        """
+
+        if driver.provider == GraphProvider.NEPTUNE:
+            query = """
+                MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
                 RETURN [x IN split(e.fact_embedding, ",") | toFloat(x)] as fact_embedding
             """
-        else:
-            query: LiteralString = """
-            MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
-            RETURN e.fact_embedding AS fact_embedding
+
+        if driver.provider == GraphProvider.KUZU:
+            query = """
+                MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(m:Entity)
+                RETURN e.fact_embedding AS fact_embedding
             """
 
         records, _, _ = await driver.execute_query(
@@ -253,15 +289,22 @@ class EntityEdge(Edge):
             'invalid_at': self.invalid_at,
         }
 
-        edge_data.update(self.attributes or {})
+        if driver.provider == GraphProvider.KUZU:
+            edge_data['attributes'] = json.dumps(self.attributes)
+            result = await driver.execute_query(
+                get_entity_edge_save_query(driver.provider),
+                **edge_data,
+            )
+        else:
+            edge_data.update(self.attributes or {})
 
-        if driver.provider == GraphProvider.NEPTUNE:
-            driver.save_to_aoss('edge_name_and_fact', [edge_data])  # pyright: ignore reportAttributeAccessIssue
+            if driver.provider == GraphProvider.NEPTUNE:
+                driver.save_to_aoss('edge_name_and_fact', [edge_data])  # pyright: ignore reportAttributeAccessIssue
 
-        result = await driver.execute_query(
-            get_entity_edge_save_query(driver.provider),
-            edge_data=edge_data,
-        )
+            result = await driver.execute_query(
+                get_entity_edge_save_query(driver.provider),
+                edge_data=edge_data,
+            )
 
         logger.debug(f'Saved edge to Graph: {self.uuid}')
 
@@ -269,21 +312,25 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
-        records, _, _ = await driver.execute_query(
-            """
+        match_query = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(m:Entity)
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
             RETURN
             """
-            + (
-                ENTITY_EDGE_RETURN_NEPTUNE
-                if driver.provider == GraphProvider.NEPTUNE
-                else ENTITY_EDGE_RETURN
-            ),
+            + get_entity_edge_return_query(driver.provider),
             uuid=uuid,
             routing_='r',
         )
 
-        edges = [get_entity_edge_from_record(record) for record in records]
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
         if len(edges) == 0:
             raise EdgeNotFoundError(uuid)
@@ -294,22 +341,26 @@ class EntityEdge(Edge):
         if len(uuids) == 0:
             return []
 
-        records, _, _ = await driver.execute_query(
-            """
+        match_query = """
             MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m:Entity)
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
             WHERE e.uuid IN $uuids
             RETURN
             """
-            + (
-                ENTITY_EDGE_RETURN_NEPTUNE
-                if driver.provider == GraphProvider.NEPTUNE
-                else ENTITY_EDGE_RETURN
-            ),
+            + get_entity_edge_return_query(driver.provider),
             uuids=uuids,
             routing_='r',
         )
 
-        edges = [get_entity_edge_from_record(record) for record in records]
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
         return edges
 
@@ -332,23 +383,27 @@ class EntityEdge(Edge):
             else ''
         )
 
-        records, _, _ = await driver.execute_query(
-            """
+        match_query = """
             MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m:Entity)
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
             WHERE e.group_id IN $group_ids
             """
             + cursor_query
             + """
             RETURN
             """
-            + (
-                ENTITY_EDGE_RETURN_NEPTUNE
-                if driver.provider == GraphProvider.NEPTUNE
-                else ENTITY_EDGE_RETURN
-            )
+            + get_entity_edge_return_query(driver.provider)
             + with_embeddings_query
             + """
-            ORDER BY e.uuid DESC 
+            ORDER BY e.uuid DESC
             """
             + limit_query,
             group_ids=group_ids,
@@ -357,7 +412,7 @@ class EntityEdge(Edge):
             routing_='r',
         )
 
-        edges = [get_entity_edge_from_record(record) for record in records]
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
         if len(edges) == 0:
             raise GroupsEdgesNotFoundError(group_ids)
@@ -365,21 +420,25 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_node_uuid(cls, driver: GraphDriver, node_uuid: str):
-        records, _, _ = await driver.execute_query(
-            """
+        match_query = """
             MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity {uuid: $node_uuid})-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m:Entity)
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
             RETURN
             """
-            + (
-                ENTITY_EDGE_RETURN_NEPTUNE
-                if driver.provider == GraphProvider.NEPTUNE
-                else ENTITY_EDGE_RETURN
-            ),
+            + get_entity_edge_return_query(driver.provider),
             node_uuid=node_uuid,
             routing_='r',
         )
 
-        edges = [get_entity_edge_from_record(record) for record in records]
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
 
         return edges
 
@@ -479,7 +538,25 @@ def get_episodic_edge_from_record(record: Any) -> EpisodicEdge:
     )
 
 
-def get_entity_edge_from_record(record: Any) -> EntityEdge:
+def get_entity_edge_from_record(record: Any, provider: GraphProvider) -> EntityEdge:
+    episodes = record['episodes']
+    if provider == GraphProvider.KUZU:
+        attributes = json.loads(record['attributes']) if record['attributes'] else {}
+    else:
+        attributes = record['attributes']
+        attributes.pop('uuid', None)
+        attributes.pop('source_node_uuid', None)
+        attributes.pop('target_node_uuid', None)
+        attributes.pop('fact', None)
+        attributes.pop('fact_embedding', None)
+        attributes.pop('name', None)
+        attributes.pop('group_id', None)
+        attributes.pop('episodes', None)
+        attributes.pop('created_at', None)
+        attributes.pop('expired_at', None)
+        attributes.pop('valid_at', None)
+        attributes.pop('invalid_at', None)
+
     edge = EntityEdge(
         uuid=record['uuid'],
         source_node_uuid=record['source_node_uuid'],
@@ -488,26 +565,14 @@ def get_entity_edge_from_record(record: Any) -> EntityEdge:
         fact_embedding=record.get('fact_embedding'),
         name=record['name'],
         group_id=record['group_id'],
-        episodes=record['episodes'],
+        episodes=episodes,
         created_at=parse_db_date(record['created_at']),  # type: ignore
         expired_at=parse_db_date(record['expired_at']),
         valid_at=parse_db_date(record['valid_at']),
         invalid_at=parse_db_date(record['invalid_at']),
-        attributes=record['attributes'],
+        attributes=attributes,
     )
 
-    edge.attributes.pop('uuid', None)
-    edge.attributes.pop('source_node_uuid', None)
-    edge.attributes.pop('target_node_uuid', None)
-    edge.attributes.pop('fact', None)
-    edge.attributes.pop('name', None)
-    edge.attributes.pop('group_id', None)
-    edge.attributes.pop('episodes', None)
-    edge.attributes.pop('created_at', None)
-    edge.attributes.pop('expired_at', None)
-    edge.attributes.pop('valid_at', None)
-    edge.attributes.pop('invalid_at', None)
-
     return edge
 
 

graphiti_core/graph_queries.py

@@ -16,6 +16,13 @@ NEO4J_TO_FALKORDB_MAPPING = {
     'episode_content': 'Episodic',
     'edge_name_and_fact': 'RELATES_TO',
 }
+# Mapping from fulltext index names to Kuzu node labels
+INDEX_TO_LABEL_KUZU_MAPPING = {
+    'node_name_and_summary': 'Entity',
+    'community_name': 'Community',
+    'episode_content': 'Episodic',
+    'edge_name_and_fact': 'RelatesToNode_',
+}
 
 
 def get_range_indices(provider: GraphProvider) -> list[LiteralString]:
@@ -35,6 +42,9 @@ def get_range_indices(provider: GraphProvider) -> list[LiteralString]:
             'CREATE INDEX FOR ()-[e:HAS_MEMBER]-() ON (e.uuid)',
         ]
 
+    if provider == GraphProvider.KUZU:
+        return []
+
     return [
         'CREATE INDEX entity_uuid IF NOT EXISTS FOR (n:Entity) ON (n.uuid)',
         'CREATE INDEX episode_uuid IF NOT EXISTS FOR (n:Episodic) ON (n.uuid)',
@@ -68,6 +78,14 @@ def get_fulltext_indices(provider: GraphProvider) -> list[LiteralString]:
             """CREATE FULLTEXT INDEX FOR ()-[e:RELATES_TO]-() ON (e.name, e.fact, e.group_id)""",
         ]
 
+    if provider == GraphProvider.KUZU:
+        return [
+            "CALL CREATE_FTS_INDEX('Episodic', 'episode_content', ['content', 'source', 'source_description']);",
+            "CALL CREATE_FTS_INDEX('Entity', 'node_name_and_summary', ['name', 'summary']);",
+            "CALL CREATE_FTS_INDEX('Community', 'community_name', ['name']);",
+            "CALL CREATE_FTS_INDEX('RelatesToNode_', 'edge_name_and_fact', ['name', 'fact']);",
+        ]
+
     return [
         """CREATE FULLTEXT INDEX episode_content IF NOT EXISTS
         FOR (e:Episodic) ON EACH [e.content, e.source, e.source_description, e.group_id]""",
@@ -80,11 +98,15 @@ def get_fulltext_indices(provider: GraphProvider) -> list[LiteralString]:
     ]
 
 
-def get_nodes_query(provider: GraphProvider, name: str = '', query: str | None = None) -> str:
+def get_nodes_query(name: str, query: str, limit: int, provider: GraphProvider) -> str:
     if provider == GraphProvider.FALKORDB:
         label = NEO4J_TO_FALKORDB_MAPPING[name]
         return f"CALL db.idx.fulltext.queryNodes('{label}', {query})"
 
+    if provider == GraphProvider.KUZU:
+        label = INDEX_TO_LABEL_KUZU_MAPPING[name]
+        return f"CALL QUERY_FTS_INDEX('{label}', '{name}', {query}, TOP := $limit)"
+
     return f'CALL db.index.fulltext.queryNodes("{name}", {query}, {{limit: $limit}})'
 
 
@@ -93,12 +115,19 @@ def get_vector_cosine_func_query(vec1, vec2, provider: GraphProvider) -> str:
         # FalkorDB uses a different syntax for regular cosine similarity and Neo4j uses normalized cosine similarity
         return f'(2 - vec.cosineDistance({vec1}, vecf32({vec2})))/2'
 
+    if provider == GraphProvider.KUZU:
+        return f'array_cosine_similarity({vec1}, {vec2})'
+
     return f'vector.similarity.cosine({vec1}, {vec2})'
 
 
-def get_relationships_query(name: str, provider: GraphProvider) -> str:
+def get_relationships_query(name: str, limit: int, provider: GraphProvider) -> str:
     if provider == GraphProvider.FALKORDB:
         label = NEO4J_TO_FALKORDB_MAPPING[name]
         return f"CALL db.idx.fulltext.queryRelationships('{label}', $query)"
 
+    if provider == GraphProvider.KUZU:
+        label = INDEX_TO_LABEL_KUZU_MAPPING[name]
+        return f"CALL QUERY_FTS_INDEX('{label}', '{name}', cast($query AS STRING), TOP := $limit)"
+
     return f'CALL db.index.fulltext.queryRelationships("{name}", $query, {{limit: $limit}})'