The returned set is bound to the calling action, so call this from the action body itself, + * not from a callback running on another thread, and obtain one per action rather than holding + * one across actions. Operating on a set that carries no binding throws rather than silently + * widening the operation to every partition key. + * * @param name the name of the memory set * @return the memory set */ @@ -39,6 +44,10 @@ public interface BaseLongTermMemory extends AutoCloseable { /** * Deletes the memory set. * + *
Unlike the set-scoped operations, this takes a name and applies to the key currently in + * scope, so it must be called from the action body rather than from another thread, and can + * target a different key than {@link MemorySet#delete} on a same-named set would. + * * @param name the name of the memory set to delete * @return true if the memory set was successfully deleted */ diff --git a/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java b/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java index 2cfdb87da..0af97ae63 100644 --- a/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java +++ b/api/src/main/java/org/apache/flink/agents/api/memory/MemorySet.java @@ -30,10 +30,18 @@ /** * Represents a long term memory set, a named collection of memory items. Acts as a thin proxy that * delegates all operations to the bound {@link BaseLongTermMemory}. + * + *
A set also carries the action context it was obtained in. Operations run on a worker thread + * when submitted through {@code durableExecuteAsync}, by which time the owning long term memory may + * already have switched to another partition key, so they take the context from the set rather than + * from it. A set must therefore be obtained per action and not reused across actions. */ public class MemorySet { private final String name; private @JsonIgnore BaseLongTermMemory ltm; + private @JsonIgnore String partitionKey; + private @JsonIgnore String observationId = ""; + private @JsonIgnore boolean observationSuppressed; @JsonCreator public MemorySet(@JsonProperty("name") String name) { @@ -100,10 +108,37 @@ public void setLtm(BaseLongTermMemory ltm) { this.ltm = ltm; } + /** + * Binds this set to the action context it was obtained in. Called on the mailbox thread when + * the set is created. + * + * @param partitionKey the partition key this set is scoped to + * @param observationId identifier for the owning action's observations + * @param observationSuppressed whether observation is suppressed for the owning action + */ + public void setActionContext( + String partitionKey, String observationId, boolean observationSuppressed) { + this.partitionKey = partitionKey; + this.observationId = observationId; + this.observationSuppressed = observationSuppressed; + } + public String getName() { return name; } + public String getPartitionKey() { + return partitionKey; + } + + public String getObservationId() { + return observationId; + } + + public boolean isObservationSuppressed() { + return observationSuppressed; + } + @Override public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; diff --git a/docs/content/docs/development/memory/long_term_memory.md b/docs/content/docs/development/memory/long_term_memory.md index 666c697b9..27dcab87b 100644 --- a/docs/content/docs/development/memory/long_term_memory.md +++ b/docs/content/docs/development/memory/long_term_memory.md @@ -179,6 +179,13 @@ public static void processEvent(Event event, RunnerContext ctx) throws Exception {{< /tabs >}} +{{< hint warning >}} +A memory set is scoped to the key of the action that obtained it. Call `get_memory_set` / +`getMemorySet` inside each action that needs one, rather than caching a set and reusing it +in a later action. Reusing a set would apply another key's operations to the key it was +originally obtained for, and operating on a set that carries no scope raises an error. +{{< /hint >}} + ### Adding Items {{< tabs "Adding Items" >}} diff --git a/python/flink_agents/api/memory/long_term_memory.py b/python/flink_agents/api/memory/long_term_memory.py index 358d820e1..e06817497 100644 --- a/python/flink_agents/api/memory/long_term_memory.py +++ b/python/flink_agents/api/memory/long_term_memory.py @@ -72,12 +72,25 @@ class MemorySetItem(BaseModel): class MemorySet(BaseModel): """Represents a long term memory set contains memory items. + A set is bound to the action context it was obtained in. Operations run on a + worker thread when submitted through ``durable_execute_async``, by which time + the owning long term memory may already have switched to another partition + key, so they take the context from the set rather than from it. A set must + therefore be obtained per action and not reused across actions. + Attributes: name: The name of this memory set. + partition_key: The partition key this set is scoped to. + observation_id: Identifier for the owning action's observations. + observation_suppressed: Whether observation is suppressed for the owning + action. """ name: str ltm: "BaseLongTermMemory" = Field(default=None, exclude=True) + partition_key: str | None = Field(default=None, exclude=True) + observation_id: str = Field(default="", exclude=True) + observation_suppressed: bool = Field(default=False, exclude=True) def add( self, @@ -150,6 +163,12 @@ class BaseLongTermMemory(ABC, BaseModel): def get_memory_set(self, name: str) -> MemorySet: """Get the memory set by name. If it does not exist, create it. + The returned set is bound to the calling action, so call this from the action + body itself, not from a callback running on another thread, and obtain one per + action rather than holding one across actions. Operating on a set that carries + no binding raises rather than silently widening the operation to every + partition key. + Args: name: The name of the memory set. @@ -161,6 +180,11 @@ def get_memory_set(self, name: str) -> MemorySet: def delete_memory_set(self, name: str) -> bool: """Delete the memory set. + Unlike the set-scoped operations, this takes a name and applies to the key + currently in scope, so it must be called from the action body rather than from + another thread, and can target a different key than ``MemorySet.delete`` on a + same-named set would. + Args: name: The name of the memory set. diff --git a/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py b/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py index db87a3470..9453d63da 100644 --- a/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py +++ b/python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py @@ -123,6 +123,23 @@ def validate_config(cls, v: Any, values: Any) -> Any: return _FlinkAgentsLlmConfig, _FlinkAgentsEmbedderConfig +def _bound_partition_key(memory_set: MemorySet) -> str: + """Return the partition key the set is scoped to. + + Mem0 ignores a falsy ``agent_id`` rather than matching on it, so an unbound set + would widen every operation to all keys sharing the job id and set name, which + for a delete means deleting another key's items. Refuse the operation instead. + """ + if memory_set.partition_key is None: + msg = ( + f"Memory set {memory_set.name!r} is not bound to a partition key. " + "Obtain it with get_memory_set inside the action that uses it, rather " + "than constructing it directly or reusing one across actions." + ) + raise ValueError(msg) + return memory_set.partition_key + + class Mem0LongTermMemory(InternalBaseLongTermMemory): """Long-Term Memory backed by Mem0. @@ -429,18 +446,33 @@ def drain_ltm_observation_records(self, key: str, observation_id: str) -> str: def get_memory_set(self, name: str) -> MemorySet: """Get the memory set by name. + The current partition key and observation context are copied onto the set + so that operations submitted to a worker thread stay scoped to the action + that obtained it. Must be called on the mailbox thread. + Args: name: The name of the memory set. Returns: The memory set. """ - return MemorySet(name=name, ltm=self) + return MemorySet( + name=name, + ltm=self, + partition_key=self.key, + observation_id=self._observation_id, + observation_suppressed=self._observation_suppressed, + ) @override def delete_memory_set(self, name: str) -> bool: """Delete a memory set and all its items. + Takes a name rather than a ``MemorySet``, so it has no bound context to read + and uses the key currently in scope. It is therefore only correct on the + mailbox thread, and deleting a whole set can target a different key than + ``MemorySet.delete`` on a set of the same name would. + Args: name: The name of the memory set. @@ -485,10 +517,10 @@ def add( Returns: List of IDs of the added memories. """ - observation_key = self.key - observation_id = self._observation_id + observation_key = _bound_partition_key(memory_set) + observation_id = memory_set.observation_id observation_enabled = ( - self._update_observation_enabled and not self._observation_suppressed + self._update_observation_enabled and not memory_set.observation_suppressed ) if isinstance(memory_items, str): memory_items = [memory_items] @@ -550,10 +582,10 @@ def get( Returns: List of memory items. """ - observation_key = self.key - observation_id = self._observation_id + observation_key = _bound_partition_key(memory_set) + observation_id = memory_set.observation_id observation_enabled = ( - self._get_observation_enabled and not self._observation_suppressed + self._get_observation_enabled and not memory_set.observation_suppressed ) if ids is not None: if isinstance(ids, str): @@ -605,10 +637,10 @@ def delete(self, memory_set: MemorySet, ids: str | List[str] | None = None) -> N memory_set: The memory set to delete from. ids: Optional ID or list of IDs. If None, deletes all items. """ - observation_key = self.key - observation_id = self._observation_id + observation_key = _bound_partition_key(memory_set) + observation_id = memory_set.observation_id observation_enabled = ( - self._update_observation_enabled and not self._observation_suppressed + self._update_observation_enabled and not memory_set.observation_suppressed ) if ids is None: self._mem0_instance.delete_all( @@ -662,10 +694,10 @@ def search( Returns: List of matching memory items. """ - observation_key = self.key - observation_id = self._observation_id + observation_key = _bound_partition_key(memory_set) + observation_id = memory_set.observation_id observation_enabled = ( - self._search_observation_enabled and not self._observation_suppressed + self._search_observation_enabled and not memory_set.observation_suppressed ) result = self._mem0_instance.search( query=query, diff --git a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py index 3f3faf4e5..58a104039 100644 --- a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_long_term_memory.py @@ -330,11 +330,12 @@ def test_switch_context(ltm) -> None: memory_set.add(items="Data for key_a") ltm.switch_context("key_b", observation_id="action-b") - # key_b should have no items in the same memory set name - items = memory_set.get() - # Items from key_a should not be visible under key_b - # (They have different agent_id scoping) - assert len(items) == 0 + # The set stays scoped to key_a, so it still reads key_a's item after the + # switch rather than following the current context. + assert len(memory_set.get()) == 1 + # A set obtained under key_b is scoped to key_b, so key_a's items in the + # same-named set are not visible through it. + assert len(ltm.get_memory_set(name="context_set").get()) == 0 # Reset context ltm.switch_context("", observation_id="action-empty") diff --git a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py index 4b83aec25..d39f2083f 100644 --- a/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py +++ b/python/flink_agents/runtime/memory/mem0/tests/test_mem0_op_recording.py @@ -24,6 +24,9 @@ from typing import Any from unittest.mock import MagicMock +import pytest + +from flink_agents.api.memory.long_term_memory import MemorySet from flink_agents.runtime.memory.internal_base_long_term_memory import ( InternalBaseLongTermMemory, ) @@ -172,16 +175,15 @@ def test_context_switch_changes_observation_owner_and_current_suppression() -> N "results": [{"event": "ADD", "id": "m1", "memory": "value"}] } ltm = _make_ltm(mem0) - memory_set = ltm.get_memory_set("prefs") ltm.switch_context( "suppressed", observation_id="suppressed-action", observation_suppressed=True ) - ltm.add(memory_set, "ignored") + ltm.add(ltm.get_memory_set("prefs"), "ignored") assert _drain(ltm, "suppressed", "suppressed-action") == [] ltm.switch_context("observed", observation_id="observed-action") - ltm.add(memory_set, "recorded") + ltm.add(ltm.get_memory_set("prefs"), "recorded") assert [record["id"] for record in _drain(ltm, "observed", "observed-action")] == [ "m1" ] @@ -196,9 +198,9 @@ def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None: mem0.get_all.return_value = {"results": []} mem0.search.return_value = {"results": []} ltm = _make_ltm(mem0) - memory_set = ltm.get_memory_set("prefs") ltm.switch_context("", observation_id="empty-action") + memory_set = ltm.get_memory_set("prefs") ltm.add(memory_set, "input") ltm.get(memory_set) ltm.search(memory_set, "query", limit=5) @@ -212,3 +214,99 @@ def test_empty_context_key_is_used_consistently_for_mem0_operations() -> None: "", "", ] + + +def test_memory_set_stays_on_its_own_key_after_the_owner_switches() -> None: + mem0 = MagicMock() + mem0.add.return_value = {"results": []} + mem0.get_all.return_value = {"results": []} + mem0.search.return_value = {"results": []} + ltm = _make_ltm(mem0) + + ltm.switch_context("owner", observation_id="owner-action") + memory_set = ltm.get_memory_set("prefs") + + ltm.switch_context("other", observation_id="other-action") + ltm.add(memory_set, "input") + ltm.get(memory_set) + ltm.search(memory_set, "query", limit=5) + ltm.delete(memory_set) + + assert mem0.add.call_args.kwargs["agent_id"] == "owner" + assert mem0.get_all.call_args.kwargs["agent_id"] == "owner" + assert mem0.search.call_args.kwargs["agent_id"] == "owner" + assert mem0.delete_all.call_args.kwargs["agent_id"] == "owner" + + +def test_observations_stay_with_the_action_that_obtained_the_set() -> None: + mem0 = MagicMock() + mem0.add.return_value = { + "results": [{"event": "ADD", "id": "m1", "memory": "value"}] + } + mem0.get_all.return_value = {"results": [{"id": "m2", "memory": "stored"}]} + mem0.search.return_value = {"results": []} + ltm = _make_ltm(mem0) + + ltm.switch_context("owner", observation_id="owner-action") + memory_set = ltm.get_memory_set("prefs") + + ltm.switch_context("other", observation_id="other-action") + ltm.add(memory_set, "input") + ltm.get(memory_set) + ltm.search(memory_set, "query", limit=5) + ltm.delete(memory_set) + + assert _drain(ltm, "other", "other-action") == [] + assert [record["op"] for record in _drain(ltm, "owner", "owner-action")] == [ + "ADD", + "GET", + "SEARCH", + "DELETE_SET", + ] + + +def test_unbound_memory_set_is_refused_rather_than_widened() -> None: + mem0 = MagicMock() + ltm = _make_ltm(mem0) + unbound = MemorySet(name="prefs", ltm=ltm) + + for op in ( + lambda: ltm.add(unbound, "input"), + lambda: ltm.get(unbound), + lambda: ltm.delete(unbound), + lambda: ltm.search(unbound, "query", limit=5), + ): + with pytest.raises(ValueError, match="not bound to a partition key"): + op() + + mem0.add.assert_not_called() + mem0.get_all.assert_not_called() + mem0.delete_all.assert_not_called() + mem0.search.assert_not_called() + + +def test_suppression_follows_the_set_not_the_current_context() -> None: + mem0 = MagicMock() + mem0.add.return_value = { + "results": [{"event": "ADD", "id": "m1", "memory": "value"}] + } + ltm = _make_ltm(mem0) + + # Obtained while suppressed, used while the current context is not: the set's + # own flag decides, so nothing is recorded. + ltm.switch_context( + "owner", observation_id="owner-action", observation_suppressed=True + ) + suppressed_set = ltm.get_memory_set("prefs") + ltm.switch_context("owner", observation_id="live-action") + ltm.add(suppressed_set, "input") + assert _drain(ltm, "owner", "owner-action") == [] + + # And the reverse: obtained unsuppressed, used while the current context is + # suppressed, so the operation is still recorded. + unsuppressed_set = ltm.get_memory_set("prefs") + ltm.switch_context( + "owner", observation_id="quiet-action", observation_suppressed=True + ) + ltm.add(unsuppressed_set, "input") + assert [record["id"] for record in _drain(ltm, "owner", "live-action")] == ["m1"] diff --git a/python/flink_agents/runtime/python_java_utils.py b/python/flink_agents/runtime/python_java_utils.py index fc1b06662..e1ecac32b 100644 --- a/python/flink_agents/runtime/python_java_utils.py +++ b/python/flink_agents/runtime/python_java_utils.py @@ -385,12 +385,23 @@ def get_long_term_memory(ctx: Any) -> Any: return ctx.long_term_memory -def to_python_memory_set(name: str) -> MemorySet: - """Build a Python ``MemorySet`` from its name. Used by the Java - ``Mem0LongTermMemory`` wrapper to forward calls into Python ``Mem0LongTermMemory``, - which expects a ``MemorySet`` instance but only reads its ``name`` field. +def to_python_memory_set( + name: str, + partition_key: str, + observation_id: str = "", + observation_suppressed: bool = False, # noqa: FBT001 +) -> MemorySet: + """Build a Python ``MemorySet`` from the fields the Java side holds. Used by the + Java ``Mem0LongTermMemory`` wrapper to forward calls into Python + ``Mem0LongTermMemory``, which reads the action context off the set rather than + off itself, so the context has to travel with each forwarded call. """ - return MemorySet(name=name) + return MemorySet( + name=name, + partition_key=partition_key, + observation_id=observation_id, + observation_suppressed=observation_suppressed, + ) def mem0_items_to_java( diff --git a/python/flink_agents/runtime/tests/test_python_java_utils.py b/python/flink_agents/runtime/tests/test_python_java_utils.py index 20dce808b..d4264d46f 100644 --- a/python/flink_agents/runtime/tests/test_python_java_utils.py +++ b/python/flink_agents/runtime/tests/test_python_java_utils.py @@ -30,6 +30,7 @@ call_embedding_with_usage, convert_to_python_key_text, get_python_tool_metadata, + to_python_memory_set, wrap_to_input_event, ) @@ -99,3 +100,12 @@ def test_convert_to_python_key_text_uses_python_str() -> None: def test_convert_to_python_key_text_does_not_unpickle_explicit_bytes() -> None: assert convert_to_python_key_text(b"N.", "explicit") == "b'N.'" assert convert_to_python_key_text(b"\x80\x04N.", "explicit") == "b'\\x80\\x04N.'" + + +def test_to_python_memory_set_carries_the_action_context() -> None: + memory_set = to_python_memory_set("prefs", "owner", "owner-action", True) + + assert memory_set.name == "prefs" + assert memory_set.partition_key == "owner" + assert memory_set.observation_id == "owner-action" + assert memory_set.observation_suppressed is True diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java index 11d95ac8d..d82f58ef0 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java @@ -46,6 +46,12 @@ public class Mem0LongTermMemory implements InteranlBaseLongTermMemory { private final PythonResourceAdapter adapter; private final PyObject pyMem0; + // Defaults mirror the Python side's own defaults, so a set obtained before any + // context switch forwards the same values Python would have used itself. + private String partitionKey = ""; + private String observationId = ""; + private boolean observationSuppressed; + public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) { this.adapter = adapter; this.pyMem0 = pyMem0; @@ -54,14 +60,20 @@ public Mem0LongTermMemory(PythonResourceAdapter adapter, PyObject pyMem0) { @Override public MemorySet getMemorySet(String name) { // Mirrors Python's `Mem0LongTermMemory.get_memory_set`: a pure factory that - // returns a new MemorySet bound to this ltm; no Python call is needed. + // returns a new MemorySet bound to this ltm; no Python call is needed. The + // current action context is copied onto the set so that operations forwarded + // from a worker thread stay scoped to the action that obtained it. MemorySet ms = new MemorySet(name); ms.setLtm(this); + ms.setActionContext(partitionKey, observationId, observationSuppressed); return ms; } @Override public boolean deleteMemorySet(String name) { + // Takes a name rather than a MemorySet, so it has no bound context and the Python + // side uses the key currently in scope. It is therefore only correct on the mailbox + // thread, and can target a different key than MemorySet.delete on a same-named set. return (Boolean) adapter.callMethod(pyMem0, "delete_memory_set", Map.of("name", name)); } @@ -149,6 +161,9 @@ public void configureObservation( @Override public void switchContext( String partitionKey, String observationId, boolean observationSuppressed) { + this.partitionKey = partitionKey; + this.observationId = observationId; + this.observationSuppressed = observationSuppressed; adapter.callMethod( pyMem0, "switch_context", @@ -173,7 +188,23 @@ public void close() { } private Object buildPyMemorySet(MemorySet memorySet) { - return adapter.invoke(TO_PYTHON_MEMORY_SET, memorySet.getName()); + // Mem0 ignores a falsy agent_id rather than matching on it, so forwarding an + // unbound set would widen the operation to every key sharing the job id and set + // name, which for a delete means deleting another key's items. + if (memorySet.getPartitionKey() == null) { + throw new IllegalStateException( + String.format( + "Memory set '%s' is not bound to a partition key. Obtain it with" + + " getMemorySet inside the action that uses it, rather than" + + " constructing it directly or reusing one across actions.", + memorySet.getName())); + } + return adapter.invoke( + TO_PYTHON_MEMORY_SET, + memorySet.getName(), + memorySet.getPartitionKey(), + memorySet.getObservationId(), + memorySet.isObservationSuppressed()); } @SuppressWarnings("unchecked") diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java index f2575be45..46b9ffe02 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java @@ -32,8 +32,10 @@ import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -49,7 +51,8 @@ public class Mem0LongTermMemoryTest { void setUp() { mocks = MockitoAnnotations.openMocks(this); ltm = new Mem0LongTermMemory(mockAdapter, mockPyMem0); - when(mockAdapter.invoke(eq("python_java_utils.to_python_memory_set"), any())) + when(mockAdapter.invoke( + eq("python_java_utils.to_python_memory_set"), any(), any(), any(), any())) .thenReturn(mockPyMemorySet); } @@ -220,4 +223,32 @@ void testSwitchContextAndCloseForward() { eq(Map.of("key", "k1", "observation_id", "observation-1"))); verify(mockAdapter).callMethod(eq(mockPyMem0), eq("close"), eq(Map.of())); } + + @Test + void testUnboundSetIsRefusedRatherThanWidened() { + MemorySet unbound = new MemorySet("notes"); + unbound.setLtm(ltm); + + assertThatThrownBy(() -> ltm.add(unbound, List.of("hello"), null)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("not bound to a partition key"); + verify(mockAdapter, never()).callMethod(eq(mockPyMem0), eq("add"), any()); + } + + @Test + void testForwardedSetCarriesTheContextItWasObtainedIn() throws Exception { + ltm.switchContext("owner", "owner-action", false); + MemorySet ms = ltm.getMemorySet("notes"); + + ltm.switchContext("other", "other-action", true); + ltm.add(ms, List.of("hello"), null); + + verify(mockAdapter) + .invoke( + eq("python_java_utils.to_python_memory_set"), + eq("notes"), + eq("owner"), + eq("owner-action"), + eq(false)); + } }