graphiti-core 0.18.0__py3-none-any.whl → 0.18.2__py3-none-any.whl

graphiti_core/graphiti.py

@@ -26,7 +26,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.driver.neo4j_driver import Neo4jDriver
-from graphiti_core.edges import EntityEdge, EpisodicEdge
+from graphiti_core.edges import CommunityEdge, EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import (
@@ -93,8 +93,11 @@ load_dotenv()
 
 class AddEpisodeResults(BaseModel):
     episode: EpisodicNode
+    episodic_edges: list[EpisodicEdge]
     nodes: list[EntityNode]
     edges: list[EntityEdge]
+    communities: list[CommunityNode]
+    community_edges: list[CommunityEdge]
 
 
 class Graphiti:
@@ -113,7 +116,7 @@ class Graphiti:
         """
         Initialize a Graphiti instance.
 
-        This constructor sets up a connection to the Neo4j database and initializes
+        This constructor sets up a connection to a graph database and initializes
         the LLM client for natural language processing tasks.
 
         Parameters
@@ -148,11 +151,11 @@ class Graphiti:
 
         Notes
         -----
-        This method establishes a connection to the Neo4j database using the provided
+        This method establishes a connection to a graph database (Neo4j by default) using the provided
         credentials. It also sets up the LLM client, either using the provided client
         or by creating a default OpenAIClient.
 
-        The default database name is set to 'neo4j'. If a different database name
+        The default database name is defined during the driver’s construction. If a different database name
         is required, it should be specified in the URI or set separately after
         initialization.
 
@@ -356,10 +359,10 @@ class Graphiti:
         group_id: str | None = None,
         uuid: str | None = None,
         update_communities: bool = False,
-        entity_types: dict[str, BaseModel] | None = None,
+        entity_types: dict[str, type[BaseModel]] | None = None,
         excluded_entity_types: list[str] | None = None,
         previous_episode_uuids: list[str] | None = None,
-        edge_types: dict[str, BaseModel] | None = None,
+        edge_types: dict[str, type[BaseModel]] | None = None,
         edge_type_map: dict[tuple[str, str], list[str]] | None = None,
     ) -> AddEpisodeResults:
         """
@@ -520,9 +523,12 @@ class Graphiti:
                 self.driver, [episode], episodic_edges, hydrated_nodes, entity_edges, self.embedder
             )
 
+            communities = []
+            community_edges = []
+
             # Update any communities
             if update_communities:
-                await semaphore_gather(
+                communities, community_edges = await semaphore_gather(
                     *[
                         update_community(self.driver, self.llm_client, self.embedder, node)
                         for node in nodes
@@ -532,7 +538,14 @@ class Graphiti:
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
 
-            return AddEpisodeResults(episode=episode, nodes=nodes, edges=entity_edges)
+            return AddEpisodeResults(
+                episode=episode,
+                episodic_edges=episodic_edges,
+                nodes=hydrated_nodes,
+                edges=entity_edges,
+                communities=communities,
+                community_edges=community_edges,
+            )
 
         except Exception as e:
             raise e
@@ -542,9 +555,9 @@ class Graphiti:
         self,
         bulk_episodes: list[RawEpisode],
         group_id: str | None = None,
-        entity_types: dict[str, BaseModel] | None = None,
+        entity_types: dict[str, type[BaseModel]] | None = None,
         excluded_entity_types: list[str] | None = None,
-        edge_types: dict[str, BaseModel] | None = None,
+        edge_types: dict[str, type[BaseModel]] | None = None,
         edge_type_map: dict[tuple[str, str], list[str]] | None = None,
     ):
         """
@@ -817,7 +830,9 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    async def build_communities(self, group_ids: list[str] | None = None) -> list[CommunityNode]:
+    async def build_communities(
+        self, group_ids: list[str] | None = None
+    ) -> tuple[list[CommunityNode], list[CommunityEdge]]:
         """
         Use a community clustering algorithm to find communities of nodes. Create community nodes summarising
         the content of these communities.
@@ -846,7 +861,7 @@ class Graphiti:
             max_coroutines=self.max_coroutines,
         )
 
-        return community_nodes
+        return community_nodes, community_edges
 
     async def search(
         self,

graphiti_core/helpers.py

@@ -28,6 +28,7 @@ from numpy._typing import NDArray
 from pydantic import BaseModel
 from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphProvider
 from graphiti_core.errors import GroupIdValidationError
 
 load_dotenv()
@@ -52,12 +53,12 @@ def parse_db_date(neo_date: neo4j_time.DateTime | str | None) -> datetime | None
     )
 
 
-def get_default_group_id(db_type: str) -> str:
+def get_default_group_id(provider: GraphProvider) -> str:
     """
     This function differentiates the default group id based on the database type.
     For most databases, the default group id is an empty string, while there are database types that require a specific default group id.
     """
-    if db_type == 'falkordb':
+    if provider == GraphProvider.FALKORDB:
         return '_'
     else:
         return ''
@@ -147,7 +148,7 @@ def validate_group_id(group_id: str) -> bool:
 
 
 def validate_excluded_entity_types(
-    excluded_entity_types: list[str] | None, entity_types: dict[str, BaseModel] | None = None
+    excluded_entity_types: list[str] | None, entity_types: dict[str, type[BaseModel]] | None = None
 ) -> bool:
     """
     Validate that excluded entity types are valid type names.

graphiti_core/models/edges/edge_db_queries.py

@@ -14,43 +14,117 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from graphiti_core.driver.driver import GraphProvider
+
 EPISODIC_EDGE_SAVE = """
-        MATCH (episode:Episodic {uuid: $episode_uuid}) 
-        MATCH (node:Entity {uuid: $entity_uuid}) 
-        MERGE (episode)-[r:MENTIONS {uuid: $uuid}]->(node)
-        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-        RETURN r.uuid AS uuid"""
+    MATCH (episode:Episodic {uuid: $episode_uuid})
+    MATCH (node:Entity {uuid: $entity_uuid})
+    MERGE (episode)-[e:MENTIONS {uuid: $uuid}]->(node)
+    SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+    RETURN e.uuid AS uuid
+"""
 
 EPISODIC_EDGE_SAVE_BULK = """
     UNWIND $episodic_edges AS edge
-    MATCH (episode:Episodic {uuid: edge.source_node_uuid}) 
-    MATCH (node:Entity {uuid: edge.target_node_uuid}) 
-    MERGE (episode)-[r:MENTIONS {uuid: edge.uuid}]->(node)
-    SET r = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
-    RETURN r.uuid AS uuid
+    MATCH (episode:Episodic {uuid: edge.source_node_uuid})
+    MATCH (node:Entity {uuid: edge.target_node_uuid})
+    MERGE (episode)-[e:MENTIONS {uuid: edge.uuid}]->(node)
+    SET e = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
+    RETURN e.uuid AS uuid
+"""
+
+EPISODIC_EDGE_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
 """
 
-ENTITY_EDGE_SAVE = """
-        MATCH (source:Entity {uuid: $edge_data.source_uuid}) 
-        MATCH (target:Entity {uuid: $edge_data.target_uuid}) 
-        MERGE (source)-[r:RELATES_TO {uuid: $edge_data.uuid}]->(target)
-        SET r = $edge_data
-        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $edge_data.fact_embedding)
-        RETURN r.uuid AS uuid"""
-
-ENTITY_EDGE_SAVE_BULK = """
-    UNWIND $entity_edges AS edge
-    MATCH (source:Entity {uuid: edge.source_node_uuid}) 
-    MATCH (target:Entity {uuid: edge.target_node_uuid}) 
-    MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
-    SET r = edge
-    WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
-    RETURN edge.uuid AS uuid
+
+def get_entity_edge_save_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        return """
+            MATCH (source:Entity {uuid: $edge_data.source_uuid})
+            MATCH (target:Entity {uuid: $edge_data.target_uuid})
+            MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+            SET e = $edge_data
+            RETURN e.uuid AS uuid
+        """
+
+    return """
+        MATCH (source:Entity {uuid: $edge_data.source_uuid})
+        MATCH (target:Entity {uuid: $edge_data.target_uuid})
+        MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+        SET e = $edge_data
+        WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)
+        RETURN e.uuid AS uuid
+    """
+
+
+def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        return """
+            UNWIND $entity_edges AS edge
+            MATCH (source:Entity {uuid: edge.source_node_uuid})
+            MATCH (target:Entity {uuid: edge.target_node_uuid})
+            MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
+            SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
+            created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at, fact_embedding: vecf32(edge.fact_embedding)}
+            WITH r, edge
+            RETURN edge.uuid AS uuid
+        """
+
+    return """
+        UNWIND $entity_edges AS edge
+        MATCH (source:Entity {uuid: edge.source_node_uuid})
+        MATCH (target:Entity {uuid: edge.target_node_uuid})
+        MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
+        SET e = edge
+        WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)
+        RETURN edge.uuid AS uuid
+    """
+
+
+ENTITY_EDGE_RETURN = """
+    e.uuid AS uuid,
+    n.uuid AS source_node_uuid,
+    m.uuid AS target_node_uuid,
+    e.group_id AS group_id,
+    e.name AS name,
+    e.fact AS fact,
+    e.episodes AS episodes,
+    e.created_at AS created_at,
+    e.expired_at AS expired_at,
+    e.valid_at AS valid_at,
+    e.invalid_at AS invalid_at,
+    properties(e) AS attributes
 """
 
-COMMUNITY_EDGE_SAVE = """
-        MATCH (community:Community {uuid: $community_uuid}) 
-        MATCH (node:Entity | Community {uuid: $entity_uuid}) 
-        MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node)
-        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-        RETURN r.uuid AS uuid"""
+
+def get_community_edge_save_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        return """
+            MATCH (community:Community {uuid: $community_uuid})
+            MATCH (node {uuid: $entity_uuid})
+            MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+            SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+            RETURN e.uuid AS uuid
+        """
+
+    return """
+        MATCH (community:Community {uuid: $community_uuid})
+        MATCH (node:Entity | Community {uuid: $entity_uuid})
+        MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+        SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+        RETURN e.uuid AS uuid
+    """
+
+
+COMMUNITY_EDGE_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
+"""

graphiti_core/models/nodes/node_db_queries.py

@@ -14,39 +14,120 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from typing import Any
+
+from graphiti_core.driver.driver import GraphProvider
+
 EPISODIC_NODE_SAVE = """
-        MERGE (n:Episodic {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content, 
-        entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
-        RETURN n.uuid AS uuid"""
+    MERGE (n:Episodic {uuid: $uuid})
+    SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
+    entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+    RETURN n.uuid AS uuid
+"""
 
 EPISODIC_NODE_SAVE_BULK = """
     UNWIND $episodes AS episode
     MERGE (n:Episodic {uuid: episode.uuid})
-    SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, 
-        source: episode.source, content: episode.content, 
+    SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description,
+        source: episode.source, content: episode.content,
     entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
     RETURN n.uuid AS uuid
 """
 
-ENTITY_NODE_SAVE = """
-        MERGE (n:Entity {uuid: $entity_data.uuid})
-        SET n:$($labels)
+EPISODIC_NODE_RETURN = """
+    e.content AS content,
+    e.created_at AS created_at,
+    e.valid_at AS valid_at,
+    e.uuid AS uuid,
+    e.name AS name,
+    e.group_id AS group_id,
+    e.source_description AS source_description,
+    e.source AS source,
+    e.entity_edges AS entity_edges
+"""
+
+
+def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
+    if provider == GraphProvider.FALKORDB:
+        return f"""
+            MERGE (n:Entity {{uuid: $entity_data.uuid}})
+            SET n:{labels}
+            SET n = $entity_data
+            RETURN n.uuid AS uuid
+        """
+
+    return f"""
+        MERGE (n:Entity {{uuid: $entity_data.uuid}})
+        SET n:{labels}
         SET n = $entity_data
         WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
-        RETURN n.uuid AS uuid"""
-
-ENTITY_NODE_SAVE_BULK = """
-    UNWIND $nodes AS node
-    MERGE (n:Entity {uuid: node.uuid})
-    SET n:$(node.labels)
-    SET n = node
-    WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
-    RETURN n.uuid AS uuid
+        RETURN n.uuid AS uuid
+    """
+
+
+def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict]) -> str | Any:
+    if provider == GraphProvider.FALKORDB:
+        queries = []
+        for node in nodes:
+            for label in node['labels']:
+                queries.append(
+                    (
+                        f"""
+                        UNWIND $nodes AS node
+                        MERGE (n:Entity {{uuid: node.uuid}})
+                        SET n:{label}
+                        SET n = node
+                        WITH n, node
+                        SET n.name_embedding = vecf32(node.name_embedding)
+                        RETURN n.uuid AS uuid
+                        """,
+                        {'nodes': [node]},
+                    )
+                )
+        return queries
+
+    return """
+        UNWIND $nodes AS node
+        MERGE (n:Entity {uuid: node.uuid})
+        SET n:$(node.labels)
+        SET n = node
+        WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
+        RETURN n.uuid AS uuid
+    """
+
+
+ENTITY_NODE_RETURN = """
+    n.uuid AS uuid,
+    n.name AS name,
+    n.group_id AS group_id,
+    n.created_at AS created_at,
+    n.summary AS summary,
+    labels(n) AS labels,
+    properties(n) AS attributes
 """
 
-COMMUNITY_NODE_SAVE = """
+
+def get_community_node_save_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.FALKORDB:
+        return """
+            MERGE (n:Community {uuid: $uuid})
+            SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at, name_embedding: $name_embedding}
+            RETURN n.uuid AS uuid
+        """
+
+    return """
         MERGE (n:Community {uuid: $uuid})
         SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
         WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
-        RETURN n.uuid AS uuid"""
+        RETURN n.uuid AS uuid
+    """
+
+
+COMMUNITY_NODE_RETURN = """
+    n.uuid AS uuid,
+    n.name AS name,
+    n.name_embedding AS name_embedding,
+    n.group_id AS group_id,
+    n.summary AS summary,
+    n.created_at AS created_at
+"""