wenjin272 commented on code in PR #887:
URL: https://github.com/apache/flink-agents/pull/887#discussion_r3664065601
##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java:
##########
@@ -235,6 +235,8 @@ private void wireLongTermMemory() {
LongTermMemoryOptions.Mem0.VECTOR_STORE.getKey()));
}
longTermMemory = new Mem0LongTermMemory(pythonResourceAdapter,
(PyObject) pyLtm);
+ // The Python runner context drains LTM observation records at action
finish.
+ pythonRunnerContext.setLongTermMemory(longTermMemory);
Review Comment:
With this wiring, could we remove the existing Python-side context switch in
`PythonActionExecutor.executePythonFunction()`?
`ActionTaskContextManager.createAndSetRunnerContext()` calls
`RunnerContextImpl.switchActionContext()` before every action execution or
coroutine resume. Now that `PythonRunnerContextImpl` holds this
`Mem0LongTermMemory` wrapper, that call
already switches the same Python Mem0 instance using the same key and also
forwards the observation flags.
The following invocation in `executePythonFunction()` therefore appears
redundant and switches the LTM context a second time for the initial Python
action:
```
interpreter.invoke(
FLINK_RUNNER_CONTEXT_SWITCH_ACTION_CONTEXT,
pythonRunnerContext,
hashOfKey);
```
Could we remove this invocation, the now-unused `hashOfKey` parameter, and
the corresponding Python helper so that LTM context switching has a single
lifecycle path?
##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -555,6 +633,27 @@ private ActionTask createActionTask(Object key, Action
action, Event event) {
}
}
+ /** Returns the logical String key for Java and PyFlink keyed streams, if
supported. */
+ @Nullable
+ private String resolveEventKeyText(Object key) {
Review Comment:
Could we avoid restricting framework memory events to String-keyed streams?
Flink guarantees that the current key is non-null, so for Java inputs we
can derive the event key text directly with:
```java
String.valueOf(key)
```
For PyFlink inputs, `getCurrentKey()` returns a single-field `Row`
produced by `KeyByKeySelector`. We should first unwrap that field:
- If an explicit `key_type` was configured, the field is already the
corresponding Java value and can use `String.valueOf(logicalKey)`.
- If no `key_type` was configured, PyFlink uses `PICKLED_BYTE_ARRAY`; the
`byte[]` should be decoded to the original Python key and then converted with
`str()`.
The current implementation accepts only logical String keys, so common
keys such as `Long` or `Integer` silently disable all framework memory events
and `AgentRunBeginEvent`. Could we instead support the string representation of
any logical Flink key and document that the event's `key` attribute is an
observational representation rather than a routing key?
##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -319,6 +383,7 @@ private void processActionTaskForKey(Object key) throws
Exception {
contextManager.createAndSetRunnerContext(
actionTask,
key,
+ resolveEventKeyText(key),
Review Comment:
Could we avoid resolving `eventKeyText` inside `processActionTaskForKey()`?
The event key is invariant for the entire run, but this method is invoked
once for every action task and coroutine/continuation resume. For PyFlink keys,
`resolveEventKeyText()` may need to unwrap the key `Row`, decode pickle data,
and invoke the Python interpreter, so repeatedly resolving the same key is
unnecessary.
Could we resolve it once in `processEvent()` when an `InputEvent` actually
begins processing, then pass it to both `tryEmitAgentRunBeginEvent()` and the
action-task processing chain?
```java
tryProcessActionTaskForKey(key, eventKeyText);
processActionTaskForKey(key, eventKeyText);
```
When additional action tasks or continuations are scheduled, they can
forward the same value. For checkpoint recovery,
`tryResumeProcessActionTasks()` can resolve it once per restored processing key
before submitting the first
mailbox task.
This would avoid repeated cross-language key decoding without introducing
an additional cache or checkpointed state.
##########
runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java:
##########
@@ -118,17 +154,44 @@ public RunnerContextImpl(
this.mailboxThreadChecker = mailboxThreadChecker;
this.agentPlan = agentPlan;
this.resourceCache = resourceCache;
+ this.memoryEventSettings =
MemoryEventSettings.from(agentPlan.getConfigData());
}
public void setLongTermMemory(InteranlBaseLongTermMemory ltm) {
this.ltm = ltm;
}
- public void switchActionContext(String actionName, MemoryContext
memoryContext, String key) {
+ public void switchActionContext(
+ String actionName,
+ MemoryContext memoryContext,
+ String ltmPartitionKey,
+ @Nullable String eventKeyText,
+ boolean observationSuppressed) {
this.actionName = actionName;
this.memoryContext = memoryContext;
+ this.ltmPartitionKey = ltmPartitionKey;
+ this.eventKeyText = eventKeyText;
+ this.observationSuppressed = observationSuppressed;
+ boolean observationAllowed = !observationSuppressed && eventKeyText !=
null;
+ boolean updateObservationEnabled =
+ observationAllowed
+ && memoryEventSettings.generate(
+ MemoryEventSettings.MemoryOp.LONG_TERM_UPDATE);
+ boolean getObservationEnabled =
+ observationAllowed
+ &&
memoryEventSettings.generate(MemoryEventSettings.MemoryOp.LONG_TERM_GET);
+ boolean searchObservationEnabled =
+ observationAllowed
+ && memoryEventSettings.generate(
+ MemoryEventSettings.MemoryOp.LONG_TERM_SEARCH);
+ this.ltmObservationEnabled =
Review Comment:
The per-operation observation settings are resolved from the `AgentPlan` and
remain constant for the lifetime of the job. With the logical key always
available, the only action-specific state here is `observationSuppressed`.
Could we configure the update/get/search settings once when initializing the
LTM, and make `switchActionContext()` only switch the key and suppression
state? This would separate job-level configuration from action-level context
switching.
##########
runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java:
##########
@@ -118,17 +154,44 @@ public RunnerContextImpl(
this.mailboxThreadChecker = mailboxThreadChecker;
this.agentPlan = agentPlan;
this.resourceCache = resourceCache;
+ this.memoryEventSettings =
MemoryEventSettings.from(agentPlan.getConfigData());
}
public void setLongTermMemory(InteranlBaseLongTermMemory ltm) {
this.ltm = ltm;
}
- public void switchActionContext(String actionName, MemoryContext
memoryContext, String key) {
+ public void switchActionContext(
+ String actionName,
+ MemoryContext memoryContext,
+ String ltmPartitionKey,
+ @Nullable String eventKeyText,
Review Comment:
Could we use a single resolved key here instead of passing both
`ltmPartitionKey` and `eventKeyText`?
Both values are derived from the same Flink key and identify the same run
context. If the logical key is resolved and converted to text once, that value
could be used for both the LTM context and framework events:
```java
switchActionContext(
actionName,
memoryContext,
contextKey,
observationSuppressed);
```
This would simplify the contract and avoid using `key.hashCode()` as a
separate LTM identity, which introduces an additional collision risk.
##########
python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py:
##########
@@ -224,7 +270,13 @@ def _create_mem0_instance(self) -> Any:
return Memory(mem0_config)
@override
- def switch_context(self, key: str) -> None:
+ def switch_context(
Review Comment:
The signature here is inconsistent with
`InternalBaseLongTermMemory.switch_context()`, which only declares the `key`
parameter. Could we keep the parent and subclass signatures aligned?
##########
python/flink_agents/runtime/memory/mem0/mem0_long_term_memory.py:
##########
@@ -233,12 +285,126 @@ def switch_context(self, key: str) -> None:
Args:
key: The new key for partition isolation.
+ update_observation_enabled: Whether mutations should be recorded.
+ get_observation_enabled: Whether gets should be recorded.
+ search_observation_enabled: Whether searches should be recorded.
"""
# 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
+ if update_observation_enabled is not None:
+ self._update_observation_enabled = update_observation_enabled
+ if get_observation_enabled is not None:
+ self._get_observation_enabled = get_observation_enabled
+ if search_observation_enabled is not None:
+ self._search_observation_enabled = search_observation_enabled
+
+ 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 or getattr(self.ctx, "_j_runner_context", None) is
None:
+ 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 or getattr(self.ctx, "_j_runner_context", None) is
None:
Review Comment:
Is this check necessary? `Mem0LongTermMemory` is created with the
`FlinkRunnerContext` constructed by `create_flink_runner_context()`, where
`_j_runner_context` is always initialized with the Java runner context.
Therefore, it cannot be `None` on this path.
--
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]
