graphiti-core 0.21.0rc2__py3-none-any.whl → 0.21.0rc3__py3-none-any.whl

graphiti_core/utils/bulk_utils.py

@@ -136,12 +136,14 @@ async def add_nodes_and_edges_bulk_tx(
         entity_data: dict[str, Any] = {
             'uuid': node.uuid,
             'name': node.name,
-            'name_embedding': node.name_embedding,
             'group_id': node.group_id,
             'summary': node.summary,
             'created_at': node.created_at,
         }
 
+        if not bool(driver.aoss_client):
+            entity_data['name_embedding'] = node.name_embedding
+
         entity_data['labels'] = list(set(node.labels + ['Entity']))
         if driver.provider == GraphProvider.KUZU:
             attributes = convert_datetimes_to_strings(node.attributes) if node.attributes else {}
@@ -161,7 +163,6 @@ async def add_nodes_and_edges_bulk_tx(
             'target_node_uuid': edge.target_node_uuid,
             'name': edge.name,
             'fact': edge.fact,
-            'fact_embedding': edge.fact_embedding,
             'group_id': edge.group_id,
             'episodes': edge.episodes,
             'created_at': edge.created_at,
@@ -170,6 +171,9 @@ async def add_nodes_and_edges_bulk_tx(
             'invalid_at': edge.invalid_at,
         }
 
+        if not bool(driver.aoss_client):
+            edge_data['fact_embedding'] = edge.fact_embedding
+
         if driver.provider == GraphProvider.KUZU:
             attributes = convert_datetimes_to_strings(edge.attributes) if edge.attributes else {}
             edge_data['attributes'] = json.dumps(attributes)
@@ -195,21 +199,29 @@ async def add_nodes_and_edges_bulk_tx(
     else:
         await tx.run(get_episode_node_save_bulk_query(driver.provider), episodes=episodes)
         await tx.run(
-            get_entity_node_save_bulk_query(driver.provider, nodes),
+            get_entity_node_save_bulk_query(
+                driver.provider, nodes, has_aoss=bool(driver.aoss_client)
+            ),
             nodes=nodes,
-            has_aoss=bool(driver.aoss_client),
         )
         await tx.run(
             get_episodic_edge_save_bulk_query(driver.provider),
             episodic_edges=[edge.model_dump() for edge in episodic_edges],
         )
         await tx.run(
-            get_entity_edge_save_bulk_query(driver.provider),
+            get_entity_edge_save_bulk_query(driver.provider, has_aoss=bool(driver.aoss_client)),
             entity_edges=edges,
-            has_aoss=bool(driver.aoss_client),
         )
 
-        if driver.aoss_client:
+        if bool(driver.aoss_client):
+            for node_data, entity_node in zip(nodes, entity_nodes, strict=True):
+                if node_data.get('uuid') == entity_node.uuid:
+                    node_data['name_embedding'] = entity_node.name_embedding
+
+            for edge_data, entity_edge in zip(edges, entity_edges, strict=True):
+                if edge_data.get('uuid') == entity_edge.uuid:
+                    edge_data['fact_embedding'] = entity_edge.fact_embedding
+
             await driver.save_to_aoss(EPISODE_INDEX_NAME, episodes)
             await driver.save_to_aoss(ENTITY_INDEX_NAME, nodes)
             await driver.save_to_aoss(ENTITY_EDGE_INDEX_NAME, edges)

{graphiti_core-0.21.0rc2.dist-info → graphiti_core-0.21.0rc3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.21.0rc2
+Version: 0.21.0rc3
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti

{graphiti_core-0.21.0rc2.dist-info → graphiti_core-0.21.0rc3.dist-info}/RECORD

@@ -64,7 +64,7 @@ graphiti_core/search/search_utils.py,sha256=HqNYbyQklAq5GoOUo_W8Xut1GMGVYlNv0kE0
 graphiti_core/telemetry/__init__.py,sha256=5kALLDlU9bb2v19CdN7qVANsJWyfnL9E60J6FFgzm3o,226
 graphiti_core/telemetry/telemetry.py,sha256=47LrzOVBCcZxsYPsnSxWFiztHoxYKKxPwyRX0hnbDGc,3230
 graphiti_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/utils/bulk_utils.py,sha256=XyjWg_S6sQ1hM8HfHHoDnDn5yAWO1Ja0fYGFMfS7XV8,17071
+graphiti_core/utils/bulk_utils.py,sha256=9XWXqjxiu2ydKMLKQRTbvzO6cO1o1HRjjpmaf5Ym51k,17633
 graphiti_core/utils/datetime_utils.py,sha256=J-zYSq7-H-2n9hYOXNIun12kM10vNX9mMATGR_egTmY,1806
 graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4wPx3Vt6aAPaXh78,301
 graphiti_core/utils/maintenance/community_operations.py,sha256=XMiokEemn96GlvjkOvbo9hIX04Fea3eVj408NHG5P4o,11042
@@ -74,7 +74,7 @@ graphiti_core/utils/maintenance/node_operations.py,sha256=r9ilkA01eq1z-nF8P_s1EX
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=IIaVtShpVkOYe6haxz3a1x3v54-MzaEXG8VsxFUNeoY,3582
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/utils/ontology_utils/entity_types_utils.py,sha256=4eVgxLWY6Q8k9cRJ5pW59IYF--U4nXZsZIGOVb_yHfQ,1285
-graphiti_core-0.21.0rc2.dist-info/METADATA,sha256=Pqw2f7ySYgCDyyuWDBvT4KLrL-SCJXa2aFHlNzIJzJo,26933
-graphiti_core-0.21.0rc2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-graphiti_core-0.21.0rc2.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.21.0rc2.dist-info/RECORD,,
+graphiti_core-0.21.0rc3.dist-info/METADATA,sha256=7a30OAvffGwDCJRWGF5Mr_AP3p-JXc1f1pQXVJcQ4RY,26933
+graphiti_core-0.21.0rc3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+graphiti_core-0.21.0rc3.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.21.0rc3.dist-info/RECORD,,