feat(episodic): persist whitelisted episode_metadata provenance keys as node properties (FalkorDB/Neo4j)

graphiti_core/driver/falkordb/operations/episode_node_ops.py

@@ -28,7 +28,7 @@
     get_episode_node_save_bulk_query,
     get_episode_node_save_query,
 )
-from graphiti_core.nodes import EpisodicNode
+from graphiti_core.nodes import EPISODIC_PROVENANCE_KEYS, EpisodicNode
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +52,12 @@ async def save(
             'valid_at': node.valid_at,
             'source': node.source.value,
         }
+        # The save query for this provider writes a full property map that
+        # includes the whitelisted provenance params (see
+        # EPISODIC_PROVENANCE_KEYS); they must always be present.
+        meta = node.episode_metadata or {}
+        for key in EPISODIC_PROVENANCE_KEYS:
+            params[key] = meta.get(key)
         if tx is not None:
             await tx.run(query, **params)
         else:
@@ -71,6 +77,11 @@ async def save_bulk(
             ep = dict(node)
             ep['source'] = str(ep['source'].value)
             ep.pop('labels', None)
+            # Same provenance contract as save(): the bulk query for this
+            # provider references the whitelisted keys on every episode.
+            meta = ep.get('episode_metadata') or {}
+            for key in EPISODIC_PROVENANCE_KEYS:
+                ep[key] = meta.get(key)
             episodes.append(ep)
 
         query = get_episode_node_save_bulk_query(GraphProvider.FALKORDB)

graphiti_core/driver/neo4j/operations/episode_node_ops.py

@@ -28,7 +28,7 @@
     get_episode_node_save_bulk_query,
     get_episode_node_save_query,
 )
-from graphiti_core.nodes import EpisodicNode
+from graphiti_core.nodes import EPISODIC_PROVENANCE_KEYS, EpisodicNode
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +52,12 @@ async def save(
             'valid_at': node.valid_at,
             'source': node.source.value,
         }
+        # The save query for this provider writes a full property map that
+        # includes the whitelisted provenance params (see
+        # EPISODIC_PROVENANCE_KEYS); they must always be present.
+        meta = node.episode_metadata or {}
+        for key in EPISODIC_PROVENANCE_KEYS:
+            params[key] = meta.get(key)
         if tx is not None:
             await tx.run(query, **params)
         else:
@@ -71,6 +77,11 @@ async def save_bulk(
             ep = dict(node)
             ep['source'] = str(ep['source'].value)
             ep.pop('labels', None)
+            # Same provenance contract as save(): the bulk query for this
+            # provider references the whitelisted keys on every episode.
+            meta = ep.get('episode_metadata') or {}
+            for key in EPISODIC_PROVENANCE_KEYS:
+                ep[key] = meta.get(key)
             episodes.append(ep)
 
         query = get_episode_node_save_bulk_query(GraphProvider.NEO4J)

graphiti_core/graphiti.py

@@ -15,6 +15,7 @@
 """
 
 import logging
+import typing
 from datetime import datetime
 from time import time
 from uuid import uuid4
@@ -995,6 +996,7 @@ async def add_episode(
         custom_extraction_instructions: str | None = None,
         saga: str | SagaNode | None = None,
         saga_previous_episode_uuid: str | None = None,
+        episode_metadata: dict[str, typing.Any] | None = None,
     ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
@@ -1108,6 +1110,7 @@ async def add_episode_endpoint(episode_data: EpisodeData):
                         source_description=source_description,
                         created_at=now,
                         valid_at=reference_time,
+                        episode_metadata=episode_metadata,
                     )
                 )
 
@@ -1328,6 +1331,7 @@ async def add_episode_bulk(
                         group_id=group_id,
                         created_at=now,
                         valid_at=episode.reference_time,
+                        episode_metadata=episode.episode_metadata,
                     )
                     for episode in bulk_episodes
                 ]

graphiti_core/models/nodes/node_db_queries.py

@@ -54,14 +54,16 @@ def get_episode_node_save_query(provider: GraphProvider) -> str:
             return """
                 MERGE (n:Episodic {uuid: $uuid})
                 SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
-                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at,
+                entry_id: $entry_id, thread_id: $thread_id, chunk_index: $chunk_index, total_chunks: $total_chunks}
                 RETURN n.uuid AS uuid
             """
         case _:  # Neo4j
             return """
                 MERGE (n:Episodic {uuid: $uuid})
                 SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
-                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at,
+                entry_id: $entry_id, thread_id: $thread_id, chunk_index: $chunk_index, total_chunks: $total_chunks}
                 RETURN n.uuid AS uuid
             """
 
@@ -96,15 +98,17 @@ def get_episode_node_save_bulk_query(provider: GraphProvider) -> str:
                 UNWIND $episodes AS episode
                 MERGE (n:Episodic {uuid: episode.uuid})
                 SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, source: episode.source, content: episode.content, 
-                entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
+                entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at,
+                entry_id: episode.entry_id, thread_id: episode.thread_id, chunk_index: episode.chunk_index, total_chunks: episode.total_chunks}
                 RETURN n.uuid AS uuid
             """
         case _:  # Neo4j
             return """
                 UNWIND $episodes AS episode
                 MERGE (n:Episodic {uuid: episode.uuid})
                 SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, source: episode.source, content: episode.content, 
-                entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
+                entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at,
+                entry_id: episode.entry_id, thread_id: episode.thread_id, chunk_index: episode.chunk_index, total_chunks: episode.total_chunks}
                 RETURN n.uuid AS uuid
             """
 

graphiti_core/nodes.py

@@ -315,6 +315,14 @@ async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]): ...
 
 
+# Whitelisted scalar episode_metadata keys persisted as first-class Episodic
+# properties on providers whose save queries write a full property map
+# (FalkorDB / Neo4j). First-class properties survive the full-map SET (a
+# post-hoc stamp would be wiped on the next re-save) and are indexable —
+# e.g. an :Episodic(entry_id) index for provenance backtraces.
+EPISODIC_PROVENANCE_KEYS = ('entry_id', 'thread_id', 'chunk_index', 'total_chunks')
+
+
 class EpisodicNode(Node):
     source: EpisodeType = Field(description='source type')
     source_description: str = Field(description='description of the data source')
@@ -350,6 +358,15 @@ async def save(self, driver: GraphDriver):
             'source': self.source.value,
         }
 
+        if driver.provider in (GraphProvider.FALKORDB, GraphProvider.NEO4J):
+            # Persist whitelisted provenance metadata as node properties
+            # (see EPISODIC_PROVENANCE_KEYS). Gated to the providers whose
+            # save queries reference the params: Kuzu has a typed schema
+            # (unknown params error) and Neptune is untested for this.
+            meta = self.episode_metadata or {}
+            for key in EPISODIC_PROVENANCE_KEYS:
+                episode_args[key] = meta.get(key)
+
         result = await driver.execute_query(
             get_episode_node_save_query(driver.provider), **episode_args
         )

graphiti_core/utils/bulk_utils.py

@@ -40,7 +40,7 @@
     get_entity_node_save_bulk_query,
     get_episode_node_save_bulk_query,
 )
-from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
+from graphiti_core.nodes import EPISODIC_PROVENANCE_KEYS, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
 from graphiti_core.utils.maintenance.dedup_helpers import (
     DedupResolutionState,
@@ -105,6 +105,7 @@ class RawEpisode(BaseModel):
     source_description: str
     source: EpisodeType
     reference_time: datetime
+    episode_metadata: dict[str, Any] | None = Field(default=None)
 
 
 async def retrieve_previous_episodes_bulk(
@@ -158,9 +159,17 @@ async def add_nodes_and_edges_bulk_tx(
     driver: GraphDriver,
 ):
     episodes = [dict(episode) for episode in episodic_nodes]
+    stamp_provenance = driver.provider in (GraphProvider.FALKORDB, GraphProvider.NEO4J)
     for episode in episodes:
         episode['source'] = str(episode['source'].value)
         episode.pop('labels', None)
+        if stamp_provenance:
+            # The bulk save queries for these providers reference the
+            # whitelisted provenance keys (see EPISODIC_PROVENANCE_KEYS);
+            # every episode dict must carry them, defaulting to null.
+            meta = episode.get('episode_metadata') or {}
+            for key in EPISODIC_PROVENANCE_KEYS:
+                episode[key] = meta.get(key)
 
     nodes = []
 

tests/test_episodic_provenance.py

@@ -0,0 +1,177 @@
+"""Unit tests for whitelisted episode_metadata provenance properties.
+
+Watercooler fork feature: EPISODIC_PROVENANCE_KEYS from episode_metadata
+are persisted as first-class Episodic node properties on FalkorDB/Neo4j
+(whose save queries write a full property map — a post-hoc stamp would be
+wiped on the next re-save). Driver-free: asserts the params passed to
+execute_query and the query text per provider.
+"""
+
+import asyncio
+from datetime import datetime, timezone
+from unittest.mock import AsyncMock, MagicMock
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.models.nodes.node_db_queries import (
+    get_episode_node_save_bulk_query,
+    get_episode_node_save_query,
+)
+from graphiti_core.nodes import EPISODIC_PROVENANCE_KEYS, EpisodeType, EpisodicNode
+
+METADATA = {
+    'entry_id': '01TESTAAAAAAAAAAAAAAAAAAAA',
+    'thread_id': 'some-topic',
+    'chunk_index': 2,
+    'total_chunks': 3,
+}
+
+
+def _node(metadata=None) -> EpisodicNode:
+    return EpisodicNode(
+        name='ep',
+        group_id='g',
+        source=EpisodeType.message,
+        source_description='desc',
+        content='body',
+        valid_at=datetime.now(timezone.utc),
+        episode_metadata=metadata,
+    )
+
+
+def _driver(provider: GraphProvider):
+    driver = MagicMock()
+    driver.provider = provider
+    driver.graph_operations_interface = None
+    driver.execute_query = AsyncMock(return_value=None)
+    return driver
+
+
+def test_falkordb_save_params_carry_provenance():
+    driver = _driver(GraphProvider.FALKORDB)
+    asyncio.run(_node(METADATA).save(driver))
+
+    _, kwargs = driver.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert kwargs[key] == METADATA[key]
+
+
+def test_falkordb_save_params_default_to_none_without_metadata():
+    driver = _driver(GraphProvider.FALKORDB)
+    asyncio.run(_node(None).save(driver))
+
+    _, kwargs = driver.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert key in kwargs and kwargs[key] is None
+
+
+def test_kuzu_save_params_do_not_carry_provenance():
+    """Kuzu has a typed schema; unknown params must not be sent."""
+    driver = _driver(GraphProvider.KUZU)
+    asyncio.run(_node(METADATA).save(driver))
+
+    _, kwargs = driver.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert key not in kwargs
+
+
+def test_save_queries_reference_provenance_fields():
+    for provider in (GraphProvider.FALKORDB, GraphProvider.NEO4J):
+        q = get_episode_node_save_query(provider)
+        bq = get_episode_node_save_bulk_query(provider)
+        for key in EPISODIC_PROVENANCE_KEYS:
+            assert f'{key}: ${key}' in q
+            assert f'{key}: episode.{key}' in bq
+
+
+def test_kuzu_and_neptune_queries_unchanged():
+    for provider in (GraphProvider.KUZU, GraphProvider.NEPTUNE):
+        q = get_episode_node_save_query(provider)
+        for key in EPISODIC_PROVENANCE_KEYS:
+            assert key not in q
+
+
+# ---------------------------------------------------------------------------
+# Operations-path coverage (PR #2 review): the per-driver
+# EpisodeNodeOperations implementations build their own params/episode
+# dicts against the same changed queries — they must carry the whitelisted
+# provenance too, or fail at runtime with missing parameters.
+# ---------------------------------------------------------------------------
+
+
+def _ops_executor():
+    executor = MagicMock()
+    executor.execute_query = AsyncMock(return_value=None)
+    return executor
+
+
+def test_falkordb_operations_save_carries_provenance():
+    from graphiti_core.driver.falkordb.operations.episode_node_ops import (
+        FalkorEpisodeNodeOperations,
+    )
+
+    executor = _ops_executor()
+    asyncio.run(FalkorEpisodeNodeOperations().save(executor, _node(METADATA)))
+
+    _, kwargs = executor.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert kwargs[key] == METADATA[key]
+
+
+def test_falkordb_operations_save_defaults_to_none_without_metadata():
+    from graphiti_core.driver.falkordb.operations.episode_node_ops import (
+        FalkorEpisodeNodeOperations,
+    )
+
+    executor = _ops_executor()
+    asyncio.run(FalkorEpisodeNodeOperations().save(executor, _node(None)))
+
+    _, kwargs = executor.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert key in kwargs and kwargs[key] is None
+
+
+def test_falkordb_operations_save_bulk_flattens_provenance():
+    from graphiti_core.driver.falkordb.operations.episode_node_ops import (
+        FalkorEpisodeNodeOperations,
+    )
+
+    executor = _ops_executor()
+    asyncio.run(
+        FalkorEpisodeNodeOperations().save_bulk(
+            executor, [_node(METADATA), _node(None)]
+        )
+    )
+
+    _, kwargs = executor.execute_query.call_args
+    episodes = kwargs['episodes']
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert episodes[0][key] == METADATA[key]
+        assert episodes[1][key] is None
+
+
+def test_neo4j_operations_save_carries_provenance():
+    from graphiti_core.driver.neo4j.operations.episode_node_ops import (
+        Neo4jEpisodeNodeOperations,
+    )
+
+    executor = _ops_executor()
+    asyncio.run(Neo4jEpisodeNodeOperations().save(executor, _node(METADATA)))
+
+    _, kwargs = executor.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert kwargs[key] == METADATA[key]
+
+
+def test_neo4j_operations_save_bulk_flattens_provenance():
+    from graphiti_core.driver.neo4j.operations.episode_node_ops import (
+        Neo4jEpisodeNodeOperations,
+    )
+
+    executor = _ops_executor()
+    asyncio.run(
+        Neo4jEpisodeNodeOperations().save_bulk(executor, [_node(METADATA)])
+    )
+
+    _, kwargs = executor.execute_query.call_args
+    for key in EPISODIC_PROVENANCE_KEYS:
+        assert kwargs['episodes'][0][key] == METADATA[key]