From e42d3ae46c08369a7ecb26e663661a746a177eeb Mon Sep 17 00:00:00 2001
From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com>
Date: Mon, 2 Dec 2024 11:47:06 -0500
Subject: [PATCH] update episode override (#219)

* update episode override

* remove unused import
---
 graphiti_core/graphiti.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py
index 74d5edec..0f428f51 100644
--- a/graphiti_core/graphiti.py
+++ b/graphiti_core/graphiti.py
@@ -318,17 +318,21 @@ async def add_episode_endpoint(episode_data: EpisodeData):
             previous_episodes = await self.retrieve_episodes(
                 reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
             )
-            episode = EpisodicNode(
-                name=name,
-                group_id=group_id,
-                labels=[],
-                source=source,
-                content=episode_body,
-                source_description=source_description,
-                created_at=now,
-                valid_at=reference_time,
+
+            episode = (
+                await EpisodicNode.get_by_uuid(self.driver, uuid)
+                if uuid is not None
+                else EpisodicNode(
+                    name=name,
+                    group_id=group_id,
+                    labels=[],
+                    source=source,
+                    content=episode_body,
+                    source_description=source_description,
+                    created_at=now,
+                    valid_at=reference_time,
+                )
             )
-            episode.uuid = uuid if uuid is not None else episode.uuid
 
             # Extract entities as nodes