graphiti-core 0.7.8__tar.gz → 0.8.0__tar.gz

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.7.8
+Version: 0.8.0
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/graphiti_core/edges.py

@@ -337,6 +337,32 @@ class EntityEdge(Edge):
             raise GroupsEdgesNotFoundError(group_ids)
         return edges
 
+    @classmethod
+    async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str):
+        query: LiteralString = """
+        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+        RETURN DISTINCT
+            e.uuid AS uuid,
+            n.uuid AS source_node_uuid,
+            m.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.fact_embedding AS fact_embedding,
+            e.episodes AS episodes,
+            e.expired_at AS expired_at,
+            e.valid_at AS valid_at,
+            e.invalid_at AS invalid_at
+        """
+        records, _, _ = await driver.execute_query(
+            query, node_uuid=node_uuid, database_=DEFAULT_DATABASE, routing_='r'
+        )
+
+        edges = [get_entity_edge_from_record(record) for record in records]
+
+        return edges
+
 
 class CommunityEdge(Edge):
     async def save(self, driver: AsyncDriver):

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/graphiti_core/graphiti.py

@@ -29,7 +29,7 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient, OpenAIClient
-from graphiti_core.nodes import CommunityNode, EntityNode, EntityType, EpisodeType, EpisodicNode
+from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
 from graphiti_core.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults
 from graphiti_core.search.search_config_recipes import (
@@ -39,7 +39,6 @@ from graphiti_core.search.search_config_recipes import (
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
-    get_communities_by_nodes,
     get_mentioned_nodes,
     get_relevant_edges,
     get_relevant_nodes,
@@ -262,7 +261,7 @@ class Graphiti:
         group_id: str = '',
         uuid: str | None = None,
         update_communities: bool = False,
-        entity_types: dict[str, EntityType] | None = None,
+        entity_types: dict[str, BaseModel] | None = None,
     ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
@@ -708,7 +707,7 @@ class Graphiti:
             bfs_origin_node_uuids,
         )
 
-    async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
+    async def get_nodes_and_edges_by_episode(self, episode_uuids: list[str]) -> SearchResults:
         episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
         edges_list = await semaphore_gather(
@@ -719,9 +718,7 @@ class Graphiti:
 
         nodes = await get_mentioned_nodes(self.driver, episodes)
 
-        communities = await get_communities_by_nodes(self.driver, nodes)
-
-        return SearchResults(edges=edges, nodes=nodes, communities=communities)
+        return SearchResults(edges=edges, nodes=nodes, communities=[])
 
     async def add_triplet(self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode):
         if source_node.name_embedding is None:

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/graphiti_core/helpers.py

@@ -28,7 +28,7 @@ load_dotenv()
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
-MAX_REFLEXION_ITERATIONS = 2
+MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 2))
 DEFAULT_PAGE_LIMIT = 20
 
 

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/graphiti_core/nodes.py

@@ -19,7 +19,7 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from enum import Enum
 from time import time
-from typing import Any, ClassVar
+from typing import Any
 from uuid import uuid4
 
 from neo4j import AsyncDriver
@@ -39,12 +39,6 @@ from graphiti_core.utils.datetime_utils import utc_now
 logger = logging.getLogger(__name__)
 
 
-class EntityType(BaseModel):
-    type_description: ClassVar[str] = Field(
-        default='', description='Description of what the entity type represents'
-    )
-
-
 class EpisodeType(Enum):
     """
     Enumeration of different types of episodes that can be processed.

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/graphiti_core/utils/maintenance/node_operations.py

@@ -18,10 +18,11 @@ import logging
 from time import time
 
 import pydantic
+from pydantic import BaseModel
 
 from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
 from graphiti_core.llm_client import LLMClient
-from graphiti_core.nodes import EntityNode, EntityType, EpisodeType, EpisodicNode
+from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.prompts import prompt_library
 from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
 from graphiti_core.prompts.extract_nodes import (
@@ -120,7 +121,7 @@ async def extract_nodes(
     llm_client: LLMClient,
     episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
-    entity_types: dict[str, EntityType] | None = None,
+    entity_types: dict[str, BaseModel] | None = None,
 ) -> list[EntityNode]:
     start = time()
     extracted_node_names: list[str] = []
@@ -156,7 +157,8 @@ async def extract_nodes(
         'previous_episodes': [ep.content for ep in previous_episodes],
         'extracted_entities': extracted_node_names,
         'entity_types': {
-            type_name: values.type_description for type_name, values in entity_types.items()
+            type_name: values.model_json_schema().get('description')
+            for type_name, values in entity_types.items()
         }
         if entity_types is not None
         else {},
@@ -262,7 +264,7 @@ async def resolve_extracted_nodes(
     existing_nodes_lists: list[list[EntityNode]],
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
-    entity_types: dict[str, EntityType] | None = None,
+    entity_types: dict[str, BaseModel] | None = None,
 ) -> tuple[list[EntityNode], dict[str, str]]:
     uuid_map: dict[str, str] = {}
     resolved_nodes: list[EntityNode] = []
@@ -295,7 +297,7 @@ async def resolve_extracted_node(
     existing_nodes: list[EntityNode],
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
-    entity_types: dict[str, EntityType] | None = None,
+    entity_types: dict[str, BaseModel] | None = None,
 ) -> tuple[EntityNode, dict[str, str]]:
     start = time()
 
@@ -330,7 +332,7 @@ async def resolve_extracted_node(
         'attributes': [],
     }
 
-    entity_type_classes: tuple[EntityType, ...] = tuple()
+    entity_type_classes: tuple[BaseModel, ...] = tuple()
     if entity_types is not None:  # type: ignore
         entity_type_classes = entity_type_classes + tuple(
             filter(

{graphiti_core-0.7.8 → graphiti_core-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.7.8"
+version = "0.8.0"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",