From fb37565112dcc300bf8fa97a0cf6238ea4acc4f5 Mon Sep 17 00:00:00 2001
From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com>
Date: Wed, 2 Apr 2025 17:19:24 -0400
Subject: [PATCH] Retrieve episodes override (#319)

* add episode override

* add test

* make format
---
 examples/podcast/podcast_runner.py |  6 ++++++
 graphiti_core/graphiti.py          | 12 ++++++++++--
 pyproject.toml                     |  2 +-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/examples/podcast/podcast_runner.py b/examples/podcast/podcast_runner.py
index 1d556b85..47509766 100644
--- a/examples/podcast/podcast_runner.py
+++ b/examples/podcast/podcast_runner.py
@@ -70,6 +70,11 @@ async def main():
     messages = parse_podcast_messages()
 
     for i, message in enumerate(messages[3:14]):
+        episodes = await client.retrieve_episodes(
+            message.actual_timestamp, 3, group_ids=['podcast']
+        )
+        episode_uuids = [episode.uuid for episode in episodes]
+
         await client.add_episode(
             name=f'Message {i}',
             episode_body=f'{message.speaker_name} ({message.role}): {message.content}',
@@ -77,6 +82,7 @@ async def main():
             source_description='Podcast Transcript',
             group_id='podcast',
             entity_types={'Person': Person},
+            previous_episode_uuids=episode_uuids,
         )
 
 
diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py
index bea74635..4f060222 100644
--- a/graphiti_core/graphiti.py
+++ b/graphiti_core/graphiti.py
@@ -262,6 +262,7 @@ class Graphiti:
         uuid: str | None = None,
         update_communities: bool = False,
         entity_types: dict[str, BaseModel] | None = None,
+        previous_episode_uuids: list[str] | None = None,
     ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
@@ -287,6 +288,9 @@ class Graphiti:
             Optional uuid of the episode.
         update_communities : bool
             Optional. Whether to update communities with new node information
+        previous_episode_uuids : list[str] | None
+            Optional.  list of episode uuids to use as the previous episodes. If this is not provided,
+            the most recent episodes by created_at date will be used.
 
         Returns
         -------
@@ -315,8 +319,12 @@ class Graphiti:
             entity_edges: list[EntityEdge] = []
             now = utc_now()
 
-            previous_episodes = await self.retrieve_episodes(
-                reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
+            previous_episodes = (
+                await self.retrieve_episodes(
+                    reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
+                )
+                if previous_episode_uuids is None
+                else await EpisodicNode.get_by_uuids(self.driver, previous_episode_uuids)
             )
 
             episode = (
diff --git a/pyproject.toml b/pyproject.toml
index ccfd4970..8c30bf7f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.8.6"
+version = "0.8.7"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",