diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py index 691fdaf9..d217d924 100644 --- a/graphiti_core/graphiti.py +++ b/graphiti_core/graphiti.py @@ -122,6 +122,11 @@ class AddBulkEpisodeResults(BaseModel): community_edges: list[CommunityEdge] +class AddTripletResults(BaseModel): + nodes: list[EntityNode] + edges: list[EntityEdge] + + class Graphiti: def __init__( self, @@ -1015,7 +1020,9 @@ class Graphiti: return SearchResults(edges=edges, nodes=nodes) - async def add_triplet(self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode): + async def add_triplet( + self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode + ) -> AddTripletResults: if source_node.name_embedding is None: await source_node.generate_name_embedding(self.embedder) if target_node.name_embedding is None: @@ -1059,6 +1066,7 @@ class Graphiti: await create_entity_node_embeddings(self.embedder, nodes) await add_nodes_and_edges_bulk(self.driver, [], [], nodes, edges, self.embedder) + return AddTripletResults(edges=edges, nodes=nodes) async def remove_episode(self, episode_uuid: str): # Find the episode to be deleted diff --git a/pyproject.toml b/pyproject.toml index d8672d62..04061eec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "graphiti-core" description = "A temporal graph building library" -version = "0.20.3" +version = "0.20.4" authors = [ { name = "Paul Paliychuk", email = "paul@getzep.com" }, { name = "Preston Rasmussen", email = "preston@getzep.com" }, diff --git a/uv.lock b/uv.lock index a001c9ea..e2a3b855 100644 --- a/uv.lock +++ b/uv.lock @@ -783,7 +783,7 @@ wheels = [ [[package]] name = "graphiti-core" -version = "0.20.3" +version = "0.20.4" source = { editable = "." } dependencies = [ { name = "diskcache" },