fix: add_episode creates the episode when a supplied uuid is not found

graphiti_core/graphiti.py

@@ -1017,7 +1017,9 @@ async def add_episode(
         group_id : str | None
             An id for the graph partition the episode is a part of.
         uuid : str | None
-            Optional uuid of the episode.
+            Optional uuid of the episode. If an episode with this uuid already
+            exists it is reprocessed; otherwise the episode is created with
+            this uuid.
         update_communities : bool
             Optional. Whether to update communities with new node information
         entity_types : dict[str, BaseModel] | None
@@ -1095,11 +1097,20 @@ async def add_episode_endpoint(episode_data: EpisodeData):
                     else await EpisodicNode.get_by_uuids(self.driver, previous_episode_uuids)
                 )
 
-                # Get or create episode
-                episode = (
-                    await EpisodicNode.get_by_uuid(self.driver, uuid)
-                    if uuid is not None
-                    else EpisodicNode(
+                # Get or create episode. A supplied uuid that already exists in the
+                # graph reprocesses that episode; a uuid not yet in the graph creates
+                # the episode under the supplied uuid (matching the MERGE-on-uuid save
+                # semantics), so callers can pre-assign episode uuids.
+                episode = None
+                if uuid is not None:
+                    try:
+                        episode = await EpisodicNode.get_by_uuid(self.driver, uuid)
+                    except NodeNotFoundError:
+                        logger.info(
+                            f'Episode {uuid} not found in graph; creating it with the supplied uuid'
+                        )
+                if episode is None:
+                    episode = EpisodicNode(
                         name=name,
                         group_id=group_id,
                         labels=[],
@@ -1109,7 +1120,8 @@ async def add_episode_endpoint(episode_data: EpisodeData):
                         created_at=now,
                         valid_at=reference_time,
                     )
-                )
+                    if uuid is not None:
+                        episode.uuid = uuid
 
                 # Create default edge type map
                 edge_type_map_default = (

mcp_server/src/graphiti_mcp_server.py

@@ -375,7 +375,10 @@ async def add_memory(
                                - 'json': For structured data
                                - 'message': For conversation-style content
         source_description (str, optional): Description of the source
-        uuid (str, optional): Optional UUID for the episode
+        uuid (str, optional): Optional UUID for the episode. If an episode with this UUID
+                                 already exists it is reprocessed; otherwise the episode is
+                                 created with this UUID, so clients can pre-assign episode
+                                 UUIDs at save time.
         reference_time (str, optional): ISO-8601 timestamp for when the described events occurred
                                  (e.g. "2025-01-15T10:30:00Z" or "2025-01-15T10:30:00+00:00"). A
                                  timezone-naive value is interpreted as UTC. Defaults to the

tests/test_add_episode_uuid.py

@@ -0,0 +1,130 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from datetime import datetime, timezone
+from unittest.mock import AsyncMock, Mock
+from uuid import uuid4
+
+import pytest
+
+from graphiti_core.cross_encoder.client import CrossEncoderClient
+from graphiti_core.embedder import EmbedderClient
+from graphiti_core.graphiti import Graphiti
+from graphiti_core.llm_client import LLMClient
+from graphiti_core.nodes import EpisodeType, EpisodicNode
+from tests.helpers_test import group_id
+
+pytest_plugins = ('pytest_asyncio',)
+
+
+@pytest.fixture
+def empty_extraction_llm_client():
+    """LLM mock that extracts nothing, so add_episode's real pipeline runs LLM-free."""
+    llm = Mock(spec=LLMClient)
+    llm.generate_response = AsyncMock(
+        return_value={'extracted_entities': [], 'edges': [], 'entity_resolutions': []}
+    )
+    return llm
+
+
+@pytest.fixture
+def null_embedder():
+    embedder = Mock(spec=EmbedderClient)
+    embedder.create = AsyncMock(return_value=[0.0] * 1024)
+    embedder.create_batch = AsyncMock(return_value=[])
+    return embedder
+
+
+@pytest.fixture
+def null_cross_encoder():
+    return Mock(spec=CrossEncoderClient)
+
+
+def make_graphiti(graph_driver, llm, embedder, cross_encoder) -> Graphiti:
+    return Graphiti(
+        graph_driver=graph_driver,
+        llm_client=llm,
+        embedder=embedder,
+        cross_encoder=cross_encoder,
+    )
+
+
+@pytest.mark.asyncio
+async def test_add_episode_with_fresh_uuid_creates_episode(
+    graph_driver, empty_extraction_llm_client, null_embedder, null_cross_encoder
+):
+    """A client-supplied uuid that does not exist yet must CREATE the episode
+    with exactly that uuid instead of raising NodeNotFoundError."""
+    graphiti = make_graphiti(
+        graph_driver, empty_extraction_llm_client, null_embedder, null_cross_encoder
+    )
+    supplied_uuid = str(uuid4())
+
+    result = await graphiti.add_episode(
+        name='fresh-uuid episode',
+        episode_body='episode created with a client-supplied uuid',
+        source_description='test',
+        reference_time=datetime.now(timezone.utc),
+        source=EpisodeType.text,
+        group_id=group_id,
+        uuid=supplied_uuid,
+    )
+
+    assert result.episode.uuid == supplied_uuid
+    stored = await EpisodicNode.get_by_uuid(graphiti.driver, supplied_uuid)
+    assert stored.uuid == supplied_uuid
+    assert stored.name == 'fresh-uuid episode'
+    assert stored.content == 'episode created with a client-supplied uuid'
+    assert stored.group_id == group_id
+
+
+@pytest.mark.asyncio
+async def test_add_episode_with_existing_uuid_reprocesses_in_place(
+    graph_driver, empty_extraction_llm_client, null_embedder, null_cross_encoder
+):
+    """Upstream semantics preserved: an existing uuid loads the stored episode
+    (its stored content wins) and does not create a duplicate node."""
+    graphiti = make_graphiti(
+        graph_driver, empty_extraction_llm_client, null_embedder, null_cross_encoder
+    )
+    now = datetime.now(timezone.utc)
+    stored = await graphiti.add_episode(
+        name='stored episode',
+        episode_body='stored content',
+        source_description='test',
+        reference_time=now,
+        source=EpisodeType.text,
+        group_id=group_id,
+    )
+    existing_uuid = stored.episode.uuid
+
+    result = await graphiti.add_episode(
+        name='ignored name',
+        episode_body='ignored body',
+        source_description='test',
+        reference_time=now,
+        source=EpisodeType.text,
+        group_id=group_id,
+        uuid=existing_uuid,
+    )
+
+    assert result.episode.uuid == existing_uuid
+    assert result.episode.content == 'stored content'
+    records, _, _ = await graphiti.driver.execute_query(
+        'MATCH (e:Episodic {uuid: $uuid}) RETURN count(e) AS n',
+        uuid=existing_uuid,
+    )
+    assert records[0]['n'] == 1