From a29c3557d30ab5016e5db59cc2e3ada98c5ca458 Mon Sep 17 00:00:00 2001
From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com>
Date: Thu, 5 Sep 2024 14:09:19 -0400
Subject: [PATCH] fix clearing name embeddings bug (#87)

fix bug
---
 graphiti_core/nodes.py               | 7 ++++++-
 graphiti_core/search/search_utils.py | 8 +++++---
 pyproject.toml                       | 2 +-
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/graphiti_core/nodes.py b/graphiti_core/nodes.py
index c2f7c83..f30d001 100644
--- a/graphiti_core/nodes.py
+++ b/graphiti_core/nodes.py
@@ -225,7 +225,8 @@ class EntityNode(Node):
         MATCH (n:Entity {uuid: $uuid})
         RETURN
             n.uuid As uuid, 
-            n.name AS name, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
             n.created_at AS created_at, 
             n.summary AS summary
         """,
@@ -239,6 +240,7 @@ class EntityNode(Node):
                 EntityNode(
                     uuid=record['uuid'],
                     name=record['name'],
+                    name_embedding=record['name_embedding'],
                     labels=['Entity'],
                     created_at=record['created_at'].to_native(),
                     summary=record['summary'],
@@ -248,3 +250,6 @@ class EntityNode(Node):
         logger.info(f'Found Node: {uuid}')
 
         return nodes[0]
+
+
+# Node helpers
diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py
index 75432a7..3f7987b 100644
--- a/graphiti_core/search/search_utils.py
+++ b/graphiti_core/search/search_utils.py
@@ -23,7 +23,8 @@ async def get_mentioned_nodes(driver: AsyncDriver, episodes: list[EpisodicNode])
         MATCH (episode:Episodic)-[:MENTIONS]->(n:Entity) WHERE episode.uuid IN $uuids
         RETURN DISTINCT
             n.uuid As uuid, 
-            n.name AS name, 
+            n.name AS name,
+            n.name_embedding AS name_embedding
             n.created_at AS created_at, 
             n.summary AS summary
         """,
@@ -37,6 +38,7 @@ async def get_mentioned_nodes(driver: AsyncDriver, episodes: list[EpisodicNode])
             EntityNode(
                 uuid=record['uuid'],
                 name=record['name'],
+                name_embedding=record['name_embedding'],
                 labels=['Entity'],
                 created_at=record['created_at'].to_native(),
                 summary=record['summary'],
@@ -221,7 +223,7 @@ async def entity_similarity_search(
                 RETURN
                     n.uuid As uuid, 
                     n.name AS name, 
-                    n.name_embeddings AS name_embedding,
+                    n.name_embedding AS name_embedding,
                     n.created_at AS created_at, 
                     n.summary AS summary
                 ORDER BY score DESC
@@ -257,7 +259,7 @@ async def entity_fulltext_search(
     RETURN
         node.uuid AS uuid, 
         node.name AS name, 
-        node.name_embeddings AS name_embedding,
+        node.name_embedding AS name_embedding,
         node.created_at AS created_at, 
         node.summary AS summary
     ORDER BY score DESC
diff --git a/pyproject.toml b/pyproject.toml
index 81d87af..dbf2d50 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.2.1"
+version = "0.2.2"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",