wenjin272 commented on code in PR #887: URL: https://github.com/apache/flink-agents/pull/887#discussion_r3719200564
########## api/src/main/java/org/apache/flink/agents/api/event/MemoryEvent.java: ########## @@ -0,0 +1,138 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.flink.agents.api.event; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.core.json.JsonWriteFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import org.apache.flink.agents.api.Event; + +import java.io.IOException; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.UUID; +import java.util.function.BiFunction; + +/** + * Base class of the memory observation events. One event per (memory scope x operation) is emitted + * at the action finish boundary; each concrete subclass pins one of the seven operation kinds as + * its {@code type}. + * + * <p>Attributes: {@code key} — the String Flink key the operation belongs to; {@code value} — the + * operation's folded JSON value map. Framework observation events are skipped for non-String keyed Review Comment: This Javadoc is stale: the runtime now derives a textual event key for Java and PyFlink keyed streams, so non-String keyed streams are no longer skipped. Please update this description and the corresponding "for a String key" wording in `AgentRunBeginEvent`. ########## python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py: ########## @@ -233,12 +293,119 @@ def switch_context(self, key: str) -> None: Args: key: The new key for partition isolation. + observation_suppressed: Whether observation is suppressed for this action. """ # Ensure Mem0 is initialized on the mailbox thread. _ = self._mem0_instance # Ensure report token usage on the mailbox thread self._report_token_metrics() self.key = key + self._observation_suppressed = observation_suppressed + + def _record_ltm_op( + self, + op: _LtmObservationOp | str, + memory_set: str, + mem_id: str | None, + value: Any, + observation_key: str, + *, + enabled: bool = True, + ) -> None: + """Buffer one LTM operation for observation. + + Args: + op: The operation kind. + memory_set: The name of the memory set operated on. + mem_id: The affected memory id, or None for whole-set ops. + value: The stored memory content, or None when not applicable. + observation_key: Partition key captured at operation entry. + enabled: Whether this operation type is configured for observation. + """ + try: + if not enabled: + return + try: + operation = _LtmObservationOp(op) + except ValueError: + logger.warning("Skipping unknown LTM observation operation %r", op) + return + self._ltm_observation_records.put( + ( + observation_key, + _LtmObservationRecord( + op=operation.value, + set=memory_set, + id=mem_id, + value=value, + ), + ) + ) + except Exception: + logger.debug("LTM observation buffering failed; skipping", exc_info=True) + + def _record_ltm_search( + self, + memory_set: str, + query: str, + hits: List[Dict[str, Any]], + observation_key: str, + *, + enabled: bool = True, + ) -> None: + """Buffer one LTM search call for observation. + + Args: + memory_set: The set searched. + query: The search query string. + hits: The ordered matched records, each with id, value, and score. + observation_key: Partition key captured at operation entry. + enabled: Whether search observation is configured. + """ + try: + if not enabled: + return + self._ltm_observation_records.put( + ( + observation_key, + _LtmObservationRecord( + op=_LtmObservationOp.SEARCH.value, + set=memory_set, + query=query, + value=hits, + ), + ) + ) + except Exception: + logger.debug("LTM observation buffering failed; skipping", exc_info=True) + + def drain_ltm_observation_records(self, key: str) -> str: + """Pop buffered LTM records for one partition key as a JSON array. + + Called from Java on the mailbox thread at action-finish flush. Records for + other partition keys are placed back in the shared queue. + + Args: + key: Partition key whose records to drain. + + Returns: + JSON array string of the drained records. + """ + records: List[_LtmObservationRecord] = [] + other_records: List[tuple[str, _LtmObservationRecord]] = [] + while True: + try: + owner_key, record = self._ltm_observation_records.get_nowait() + except queue.Empty: + break + if owner_key == key: Review Comment: **[P2] LTM observations are scoped only by partition key.** For two async actions on the same key, action A can record an LTM operation and suspend, then action B can finish before A's continuation and drain A's record here. The event is then attributed to B; if A later fails, this also violates the contract that failed actions emit no memory events. Please add an action-scoped observation ID that is preserved across continuations and use `(key, observation ID)` for recording/draining, with an interleaving regression test. ########## docs/content/docs/development/memory/long_term_memory.md: ########## @@ -539,13 +539,15 @@ public static void processEvent(Event event, RunnerContext ctx) throws Exception ## Context Isolation -Long-Term Memory automatically provides context isolation through Flink's keyed partition model. Each keyed partition maintains its own isolated set of memories, ensuring that memories from one user or session do not leak into another. +Long-Term Memory provides context isolation through Flink's keyed partition model. When each logical key has a stable and unique textual representation, each keyed partition uses an isolated memory context so memories from one user or session do not leak into another. The isolation hierarchy works as follows: - **Job-level** (`JOB_IDENTIFIER`): Separates memories between different Flink jobs - **Partition-level** (keyed partition key): Separates memories between different keys within the same job - **Set-level** (memory set name): Separates memories between different logical categories within the same partition -This means you can reuse the same memory set name across different partitions, and each partition will normally access only its own memories. +This means you can reuse the same memory set name across different partitions, and each partition will normally access only its own memories. The partition context uses the textual representation of the logical Flink key: Java and explicitly typed PyFlink keys normally use `String.valueOf`; default pickled PyFlink keys are deserialized and use Python `str`. Explicit PyFlink primitive byte-array keys use Python's bytes representation directly, without pickle deserialization. -> **Note:** Partition-level isolation is currently derived from the hash of the partition key (`String.valueOf(key.hashCode())`) rather than the full original key. Distinct keys whose hashes collide may therefore share the same memory context. Avoid relying on isolation as a strict security boundary; if collision-free isolation is required, encode a unique identifier into the memory set name. +Keys used for Long-Term Memory should therefore have a stable and unique textual representation within the job. Distinct logical keys that produce the same text share one Mem0 context, so this representation is not a security boundary. + +> **Compatibility note:** Earlier experimental versions used `String.valueOf(key.hashCode())` as the Mem0 partition identity. Existing records stored under that hash-derived identity are not migrated automatically and are not visible through the new logical-key identity. Migrate those records explicitly before upgrading if they must remain accessible. Review Comment: Could we keep the surrounding user-facing isolation description unchanged and limit this change to updating the previous Note? For example: > **Note:** Partition-level isolation uses a textual identity derived from the logical key instead of `key.hashCode()`. Java keys use `String.valueOf(key)`. Default-serialized PyFlink keys are deserialized and use Python `str`; explicitly typed PyFlink keys use `String.valueOf`, except byte-array keys, which use Python's bytes representation. This avoids hash collisions, but distinct keys may still share a memory context if they produce the same text, for example when custom key types have non-unique `toString()` or `__str__()` implementations. Key representations should therefore be stable and unique, and this isolation should not be treated as a security boundary. The migration detail would fit better in release notes; also, the current statement that old records "are not visible" is too absolute because the old and new identities can occasionally be identical. ########## docs/content/docs/development/memory/sensory_and_short_term_memory.md: ########## @@ -291,9 +291,13 @@ Short-term memory can be configured with a time-to-live (TTL) so that older stat Set `short-term-memory.state-ttl.ms` to a value greater than 0 in milliseconds to enable TTL. You can also configure how the TTL is refreshed and whether expired state can be returned before Flink cleans it up: -- `short-term-memory.state-ttl.update-type`: controls whether TTL is refreshed on create/write or on read/write. +- `short-term-memory.state-ttl.update-type`: controls whether TTL is refreshed on create/write (`ON_CREATE_AND_WRITE`) or on read/write (`ON_READ_AND_WRITE`, the default). - `short-term-memory.state-ttl.visibility`: controls whether expired memory is never returned or may be returned if it has not been cleaned up yet. +{{< hint warning >}} +The default `ON_READ_AND_WRITE` update type extends an entry's lifetime whenever it is read. This also applies when producing the run-begin snapshot used by [Memory Events]({{< ref "docs/development/memory/memory_events" >}}): if you opt in through `agent-run.begin-event`, each input scans the key's short-term memory and refreshes TTL for the entries it reads, although only value nodes are included in the event. Choose `ON_CREATE_AND_WRITE` when entries should expire based only on writes. Review Comment: The last sentence skips the direct option of leaving the run-begin event disabled. Could we make the choices explicit? > With the default `ON_READ_AND_WRITE` update type, every read refreshes an entry's TTL. Enabling `agent-run.begin-event` introduces an additional source of reads: each input scans the key's short-term memory to produce the run-begin snapshot, which may extend the lifetime of the scanned entries even though only value nodes are included in the event. Leave `agent-run.begin-event` disabled if the snapshot is not needed. If the snapshot is needed but reads should not extend TTL, use `ON_CREATE_AND_WRITE`. The corresponding warning in `memory_events.md` should be aligned as well. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
