graphiti-core 0.17.4__py3-none-any.whl → 0.24.3__py3-none-any.whl

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
@@ -24,35 +25,22 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.embedder import EmbedderClient
 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_SAVE,
-    ENTITY_EDGE_SAVE,
+    COMMUNITY_EDGE_RETURN,
+    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
 
 logger = logging.getLogger(__name__)
 
-ENTITY_EDGE_RETURN: LiteralString = """
-        RETURN
-            e.uuid AS uuid,
-            startNode(e).uuid AS source_node_uuid,
-            endNode(e).uuid AS target_node_uuid,
-            e.created_at AS created_at,
-            e.name AS name,
-            e.group_id AS group_id,
-            e.fact AS fact,
-            e.episodes AS episodes,
-            e.expired_at AS expired_at,
-            e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at,
-            properties(e) AS attributes
-            """
-
 
 class Edge(BaseModel, ABC):
     uuid: str = Field(default_factory=lambda: str(uuid4()))
@@ -65,17 +53,68 @@ 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.graph_operations_interface:
+            return await driver.graph_operations_interface.edge_delete(self, driver)
+
+        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]):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.edge_delete_by_uuids(cls, driver, 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}')
 
     def __hash__(self):
         return hash(self.uuid)
@@ -108,14 +147,10 @@ class EpisodicEdge(Edge):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -130,15 +165,11 @@ class EpisodicEdge(Edge):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
-        WHERE e.uuid IN $uuids
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+            WHERE e.uuid IN $uuids
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -162,19 +193,17 @@ class EpisodicEdge(Edge):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
-        WHERE e.group_id IN $group_ids
-        """
+            MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
+            WHERE e.group_id IN $group_ids
+            """
             + cursor_query
             + """
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        ORDER BY e.uuid DESC 
-        """
+            RETURN
+            """
+            + EPISODIC_EDGE_RETURN
+            + """
+            ORDER BY e.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -222,11 +251,31 @@ class EntityEdge(Edge):
         return self.fact_embedding
 
     async def load_fact_embedding(self, driver: GraphDriver):
-        query: LiteralString = """
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.edge_load_embeddings(self, driver)
+
+        query = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
             RETURN e.fact_embedding AS fact_embedding
         """
-        records, _, _ = await driver.execute_query(query, uuid=self.uuid, routing_='r')
+
+        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
+            """
+
+        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(
+            query,
+            uuid=self.uuid,
+            routing_='r',
+        )
 
         if len(records) == 0:
             raise EdgeNotFoundError(self.uuid)
@@ -249,12 +298,18 @@ class EntityEdge(Edge):
             'invalid_at': self.invalid_at,
         }
 
-        edge_data.update(self.attributes or {})
-
-        result = await driver.execute_query(
-            ENTITY_EDGE_SAVE,
-            edge_data=edge_data,
-        )
+        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 {})
+            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}')
 
@@ -262,37 +317,84 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
+        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
             """
-        MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
-        """
-            + 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)
         return edges[0]
 
+    @classmethod
+    async def get_between_nodes(
+        cls, driver: GraphDriver, source_node_uuid: str, target_node_uuid: str
+    ):
+        match_query = """
+            MATCH (n:Entity {uuid: $source_node_uuid})-[e:RELATES_TO]->(m:Entity {uuid: $target_node_uuid})
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity {uuid: $source_node_uuid})
+                      -[:RELATES_TO]->(e:RelatesToNode_)
+                      -[:RELATES_TO]->(m:Entity {uuid: $target_node_uuid})
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider),
+            source_node_uuid=source_node_uuid,
+            target_node_uuid=target_node_uuid,
+            routing_='r',
+        )
+
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
+
+        return edges
+
     @classmethod
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         if len(uuids) == 0:
             return []
 
+        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
             """
-        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-        WHERE e.uuid IN $uuids
-        """
-            + 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
 
@@ -303,20 +405,40 @@ class EntityEdge(Edge):
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
+        with_embeddings: bool = False,
     ):
         cursor_query: LiteralString = 'AND e.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
+        with_embeddings_query: LiteralString = (
+            """,
+                e.fact_embedding AS fact_embedding
+                """
+            if with_embeddings
+            else ''
+        )
+
+        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
             """
-        MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
-        WHERE e.group_id IN $group_ids
-        """
             + cursor_query
-            + ENTITY_EDGE_RETURN
             + """
-        ORDER BY e.uuid DESC 
-        """
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider)
+            + with_embeddings_query
+            + """
+            ORDER BY e.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -324,7 +446,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)
@@ -332,15 +454,25 @@ class EntityEdge(Edge):
 
     @classmethod
     async def get_by_node_uuid(cls, driver: GraphDriver, node_uuid: str):
-        query: LiteralString = (
+        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
             """
-                                                        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
-                                                        """
-            + ENTITY_EDGE_RETURN
+            + get_entity_edge_return_query(driver.provider),
+            node_uuid=node_uuid,
+            routing_='r',
         )
-        records, _, _ = await driver.execute_query(query, 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
 
@@ -348,7 +480,7 @@ class EntityEdge(Edge):
 class CommunityEdge(Edge):
     async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
-            COMMUNITY_EDGE_SAVE,
+            get_community_edge_save_query(driver.provider),
             community_uuid=self.source_node_uuid,
             entity_uuid=self.target_node_uuid,
             uuid=self.uuid,
@@ -364,14 +496,10 @@ class CommunityEdge(Edge):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m:Entity | Community)
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m)
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN,
             uuid=uuid,
             routing_='r',
         )
@@ -384,15 +512,11 @@ class CommunityEdge(Edge):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
-        WHERE e.uuid IN $uuids
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        """,
+            MATCH (n:Community)-[e:HAS_MEMBER]->(m)
+            WHERE e.uuid IN $uuids
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN,
             uuids=uuids,
             routing_='r',
         )
@@ -414,19 +538,17 @@ class CommunityEdge(Edge):
 
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
-        WHERE e.group_id IN $group_ids
-        """
+            MATCH (n:Community)-[e:HAS_MEMBER]->(m)
+            WHERE e.group_id IN $group_ids
+            """
             + cursor_query
             + """
-        RETURN
-            e.uuid As uuid,
-            e.group_id AS group_id,
-            n.uuid AS source_node_uuid, 
-            m.uuid AS target_node_uuid, 
-            e.created_at AS created_at
-        ORDER BY e.uuid DESC
-        """
+            RETURN
+            """
+            + COMMUNITY_EDGE_RETURN
+            + """
+            ORDER BY e.uuid DESC
+            """
             + limit_query,
             group_ids=group_ids,
             uuid=uuid_cursor,
@@ -450,34 +572,41 @@ 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'],
         target_node_uuid=record['target_node_uuid'],
         fact=record['fact'],
+        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
 
 
@@ -492,8 +621,11 @@ def get_community_edge_from_record(record: Any):
 
 
 async def create_entity_edge_embeddings(embedder: EmbedderClient, edges: list[EntityEdge]):
-    if len(edges) == 0:
+    # filter out falsey values from edges
+    filtered_edges = [edge for edge in edges if edge.fact]
+
+    if len(filtered_edges) == 0:
         return
-    fact_embeddings = await embedder.create_batch([edge.fact for edge in edges])
-    for edge, fact_embedding in zip(edges, fact_embeddings, strict=True):
+    fact_embeddings = await embedder.create_batch([edge.fact for edge in filtered_edges])
+    for edge, fact_embedding in zip(filtered_edges, fact_embeddings, strict=True):
         edge.fact_embedding = fact_embedding

graphiti_core/embedder/azure_openai.py

@@ -17,7 +17,7 @@ limitations under the License.
 import logging
 from typing import Any
 
-from openai import AsyncAzureOpenAI
+from openai import AsyncAzureOpenAI, AsyncOpenAI
 
 from .client import EmbedderClient
 
@@ -25,9 +25,16 @@ logger = logging.getLogger(__name__)
 
 
 class AzureOpenAIEmbedderClient(EmbedderClient):
-    """Wrapper class for AsyncAzureOpenAI that implements the EmbedderClient interface."""
+    """Wrapper class for Azure OpenAI that implements the EmbedderClient interface.
 
-    def __init__(self, azure_client: AsyncAzureOpenAI, model: str = 'text-embedding-3-small'):
+    Supports both AsyncAzureOpenAI and AsyncOpenAI (with Azure v1 API endpoint).
+    """
+
+    def __init__(
+        self,
+        azure_client: AsyncAzureOpenAI | AsyncOpenAI,
+        model: str = 'text-embedding-3-small',
+    ):
         self.azure_client = azure_client
         self.model = model
 

graphiti_core/embedder/client.py

@@ -14,12 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 
 from pydantic import BaseModel, Field
 
-EMBEDDING_DIM = 1024
+EMBEDDING_DIM = int(os.getenv('EMBEDDING_DIM', 1024))
 
 
 class EmbedderConfig(BaseModel):