add triplet results

This commit is contained in:
prestonrasmussen 2025-09-08 15:16:42 -04:00
parent 469703e180
commit 7ebc5123a5
2 changed files with 68 additions and 62 deletions

View file

@ -122,6 +122,11 @@ class AddBulkEpisodeResults(BaseModel):
community_edges: list[CommunityEdge] community_edges: list[CommunityEdge]
class AddTripletResults(BaseModel):
nodes: list[EntityNode]
edges: list[EntityEdge]
class Graphiti: class Graphiti:
def __init__( def __init__(
self, self,
@ -1015,8 +1020,9 @@ class Graphiti:
return SearchResults(edges=edges, nodes=nodes) return SearchResults(edges=edges, nodes=nodes)
async def add_triplet(self, source_node: EntityNode, edge: EntityEdge, async def add_triplet(
target_node: EntityNode) -> AddEpisodeResults: self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode
) -> AddTripletResults:
if source_node.name_embedding is None: if source_node.name_embedding is None:
await source_node.generate_name_embedding(self.embedder) await source_node.generate_name_embedding(self.embedder)
if target_node.name_embedding is None: if target_node.name_embedding is None:
@ -1060,7 +1066,7 @@ class Graphiti:
await create_entity_node_embeddings(self.embedder, nodes) await create_entity_node_embeddings(self.embedder, nodes)
await add_nodes_and_edges_bulk(self.driver, [], [], nodes, edges, self.embedder) await add_nodes_and_edges_bulk(self.driver, [], [], nodes, edges, self.embedder)
return AddEpisodeResults(edges=edges, nodes=nodes) return AddTripletResults(episode=None, edges=edges, nodes=nodes)
async def remove_episode(self, episode_uuid: str): async def remove_episode(self, episode_uuid: str):
# Find the episode to be deleted # Find the episode to be deleted

2
uv.lock generated
View file

@ -783,7 +783,7 @@ wheels = [
[[package]] [[package]]
name = "graphiti-core" name = "graphiti-core"
version = "0.20.3" version = "0.20.4"
source = { editable = "." } source = { editable = "." }
dependencies = [ dependencies = [
{ name = "diskcache" }, { name = "diskcache" },