From 1e0700de3f4869a79f946f6beb507e8cd8d0e268 Mon Sep 17 00:00:00 2001 From: calebjacksonhoward Date: Wed, 8 Jul 2026 03:16:39 -0700 Subject: [PATCH 1/2] feat(episodic): persist whitelisted episode_metadata provenance keys as node properties (FalkorDB/Neo4j) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Watercooler A1-full (thread audit-transport-modes-hosted-db-2026-07, completion sequence Wave 4): entry_id/thread_id/chunk_index/total_chunks from episode_metadata become first-class Episodic properties so they survive the providers' full-map 'SET n = {...}' writes (a post-hoc stamp is wiped on the next re-save) and are indexable (:Episodic(entry_id) for provenance backtraces). - add_episode / RawEpisode gain episode_metadata passthrough into EpisodicNode (the field existed on the model but was never plumbed). - EpisodicNode.save flattens EPISODIC_PROVENANCE_KEYS into the save params, gated to FalkorDB/Neo4j (Kuzu has a typed schema — unknown params error; Neptune untested — both skipped with in-code notes). - Single + bulk save queries for FalkorDB/Neo4j write the four fields; Kuzu/Neptune queries unchanged. - bulk path: add_nodes_and_edges_bulk_tx flattens per episode dict, same provider gate. Tests: driver-free unit tests (params carry/omit per provider; query text references; Kuzu/Neptune untouched) — dependency-free via asyncio.run so they pass with or without pytest-asyncio. Live-validated against a real FalkorDB: properties persist, survive full-map re-save, and default to null without metadata. Co-Authored-By: Claude Fable 5 Signed-off-by: calebjacksonhoward --- graphiti_core/graphiti.py | 4 + graphiti_core/models/nodes/node_db_queries.py | 12 ++- graphiti_core/nodes.py | 17 ++++ graphiti_core/utils/bulk_utils.py | 11 ++- tests/test_episodic_provenance.py | 90 +++++++++++++++++++ 5 files changed, 129 insertions(+), 5 deletions(-) create mode 100644 tests/test_episodic_provenance.py diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py index 51c3099fa4..9781f526d8 100644 --- a/graphiti_core/graphiti.py +++ b/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 ] diff --git a/graphiti_core/models/nodes/node_db_queries.py b/graphiti_core/models/nodes/node_db_queries.py index 3c13253ccb..e5590d9ff7 100644 --- a/graphiti_core/models/nodes/node_db_queries.py +++ b/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,7 +98,8 @@ 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 @@ -104,7 +107,8 @@ 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 """ diff --git a/graphiti_core/nodes.py b/graphiti_core/nodes.py index 9644ec2b55..f1b4d839d6 100644 --- a/graphiti_core/nodes.py +++ b/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 ) diff --git a/graphiti_core/utils/bulk_utils.py b/graphiti_core/utils/bulk_utils.py index 92f25582e2..5fcdbd539e 100644 --- a/graphiti_core/utils/bulk_utils.py +++ b/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 = [] diff --git a/tests/test_episodic_provenance.py b/tests/test_episodic_provenance.py new file mode 100644 index 0000000000..acc71f502c --- /dev/null +++ b/tests/test_episodic_provenance.py @@ -0,0 +1,90 @@ +"""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 From eb3bfc4080b42a3c2ea261821c60cee986f25953 Mon Sep 17 00:00:00 2001 From: calebjacksonhoward Date: Wed, 8 Jul 2026 14:49:04 -0700 Subject: [PATCH 2/2] fix(operations): carry provenance params through the per-driver EpisodeNodeOperations save paths (PR #2 review) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The review caught that the FalkorDB/Neo4j operations-interface save() built its own params dict against the changed full-map query — missing the now-required provenance params (runtime failure) — and save_bulk() serialized dict(node) without flattening episode_metadata (silent provenance drop). Both now apply the same EPISODIC_PROVENANCE_KEYS flattening/defaulting as EpisodicNode.save() and add_nodes_and_edges_bulk_tx(). Five operations-path tests added (save carry/default, bulk flatten, both providers); note the plain FalkorDB driver does not activate the operations interface, so the classic path remains the one our deployment exercises — this closes the gap for callers that do enable it. Co-Authored-By: Claude Fable 5 Signed-off-by: calebjacksonhoward --- .../falkordb/operations/episode_node_ops.py | 13 ++- .../neo4j/operations/episode_node_ops.py | 13 ++- tests/test_episodic_provenance.py | 87 +++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/graphiti_core/driver/falkordb/operations/episode_node_ops.py b/graphiti_core/driver/falkordb/operations/episode_node_ops.py index a6ef6d48f0..af92aa70c7 100644 --- a/graphiti_core/driver/falkordb/operations/episode_node_ops.py +++ b/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) diff --git a/graphiti_core/driver/neo4j/operations/episode_node_ops.py b/graphiti_core/driver/neo4j/operations/episode_node_ops.py index 2efbf3aa63..068f78ac39 100644 --- a/graphiti_core/driver/neo4j/operations/episode_node_ops.py +++ b/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) diff --git a/tests/test_episodic_provenance.py b/tests/test_episodic_provenance.py index acc71f502c..6b1d5c9a5b 100644 --- a/tests/test_episodic_provenance.py +++ b/tests/test_episodic_provenance.py @@ -88,3 +88,90 @@ def test_kuzu_and_neptune_queries_unchanged(): 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]